tests/jax/test_distributed_fused_attn.py · 7669bf3da68074517b134cd6acebd04b221fd545 · OpenDAS / TransformerEngine

Kirthi Shankar Sivamani authored Jun 13, 2024



* Apply formatting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Apply formatting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

9416519d

test_distributed_fused_attn.py 10.7 KB

Replace test_distributed_fused_attn.py