transformer_engine/jax/cpp_extensions/attention.py · b0d562d8ac3f0ce36131471ae03d87b90a797e6f · OpenDAS / TransformerEngine

[JAX] Fix `rng_state` shape in fused attention (#2217) · b0d562d8

Phuong Nguyen authored Oct 01, 2025



fix rng_state shape
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>

b0d562d8

attention.py 107 KB

Replace attention.py