transformer_engine/jax/cpp_extensions/attention.py · bee4649c15a79ffcb9689ca7c0c963f5febaa28a · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Fix softmax aux shapes for packed/THD format (#1575) · bee4649c
Michael Goldfarb authored Mar 18, 2025
```
* Fix softmax shape for THD format.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
```
bee4649c

attention.py 102 KB