transformer_engine/jax/layernorm_mlp.py · 397c4be6629d9e6345355c623e9a75c5c59a5356 · OpenDAS / TransformerEngine

[JAX] TensorUsage + FP8 GEMM with all layouts handling on BW (#1844) · 3a298e6b

Phuong Nguyen authored Jun 18, 2025



* TensorUsage + FP8 GEMM with all layouts handling on BW
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


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

3a298e6b

layernorm_mlp.py 15.5 KB

Replace layernorm_mlp.py