tests/jax/test_distributed_fused_attn.py · d3cbccdf9e98a5b3eb756a61e5c1a744b6daf06f · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Scale sequence length in CP tests to avoid tiny sizes. (#1347) · d3cbccdf
Michael Goldfarb authored Dec 04, 2024
```
Scale sequence length in CP tests to avoid tiny sizes.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
```
d3cbccdf

test_distributed_fused_attn.py 23 KB