tests/jax/test_distributed_fused_attn.py · c42beef4e1d13f8d99e2892e875cb9e96c2d67f7 · OpenDAS / TransformerEngine

[JAX] Fix correctness of JAX fused attention with CP and improve numerics... · 20c75295

Michael Goldfarb authored Oct 24, 2024


[JAX] Fix correctness of JAX fused attention with CP and improve numerics check in unit tests (#1282)

Fix correctness of JAX fused attention with CP.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>

20c75295

test_distributed_fused_attn.py 20.9 KB

Replace test_distributed_fused_attn.py