tests/jax/test_fused_attn.py · d99142a0177a2462cbda07a31aaa8e68b4e85461 · OpenDAS / TransformerEngine

[JAX] Splitting cpp_extensions.py (#899) · 5986342a

Phuong Nguyen authored Jun 13, 2024



* Splitted cpp_extensions.py, renamed mlp.py and fused_attn.py
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* fixed import in tests
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

5986342a

test_fused_attn.py 19.6 KB

Replace test_fused_attn.py