transformer_engine/jax/cpp_extensions/attention.py · 20c55e4618138ab6398a4e5984dd3afc3aa60ae5 · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

Check for backend support in Jax context parallel fused attention test (#1227) · 20c55e46
Michael Goldfarb authored Oct 15, 2024
```
Update test to check support for context parallel attention.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
```
20c55e46

attention.py 58.2 KB