transformer_engine/jax/layernorm_dense.py · 3a298e6b194c6cc659f79533a49a13b26cd9f560 · 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_dense.py 9.01 KB

Replace layernorm_dense.py