transformer_engine/jax/cpp_extensions/gemm.py · d972e76d0688b7f3441df15fc6e36555106c4817 · 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