transformer_engine/jax/cpp_extensions/gemm.py · 557f0cb59859947e334dc33e1bf8ac488e207dcb · OpenDAS / TransformerEngine

[JAX] Removes unneccessary reshapes for FP8 GEMM (#1820) · c9e8e305

Phuong Nguyen authored May 28, 2025



* removes unnecessary reshapes for FP8 GEMM

* use nn.jax.scaled_matmul
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

c9e8e305

gemm.py 17.1 KB

Replace gemm.py