transformer_engine/jax/cpp_extensions/gemm.py · 2e23ad7127217728b55f6f11a1ed12638ee20a8a · OpenDAS / TransformerEngine

Phuong Nguyen authored Aug 25, 2025



* added shardy warning
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


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

2e23ad71

gemm.py 52.6 KB