Open-source fused Gated DeltaNet-2 training kernels for TPU v5e in JAX/Pallas

Article automatically generated from technical news.

I released an independent open-source implementation of fused Gated DeltaNet-2 training kernels for TPU v5e using JAX/Pallas. Repository: https://github.com/Akseleu-J/atomic-ops Citable release: https://doi.org/10.5281/zenodo.22706659 The repository contains three implementations I wrote for the same computation: - an earlier `jax.lax.associative_scan` path; - a pure-JAX chunked-WY path; - a fused TPU v5e Pallas path. On TPU v5e-8, for B=8, L=4096, 6 head