tests/jax/test_distributed_layernorm_mlp.py · 5986342a8fbef9ea3804542c6a91859528a2b333 · 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_distributed_layernorm_mlp.py 13.9 KB

Replace test_distributed_layernorm_mlp.py