[JAX] Splitting cpp_extensions.py (#899)

* 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>

[JAX] Splitting cpp_extensions.py (#899)
* 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 · Phuong Nguyen · GitHub · b5a7c9f9 · 5986342a · 5986342a
Unverified Commit 5986342a authored Jun 13, 2024 by Phuong Nguyen Committed by GitHub Jun 13, 2024
4 changed files
--- a/transformer_engine/jax/layernorm.py
+++ b/transformer_engine/jax/layernorm.py
--- a/transformer_engine/jax/mlp.py
+++ b/transformer_engine/jax/mlp.py
--- a/transformer_engine/jax/praxis/transformer.py
+++ b/transformer_engine/jax/praxis/transformer.py
--- a/transformer_engine/jax/softmax.py
+++ b/transformer_engine/jax/softmax.py