transformer_engine/jax/sharding.py · ed42b5ac6fd0e6cf48ef037adcb875f268d94151 · OpenDAS / TransformerEngine

Phuong Nguyen authored Aug 06, 2025



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

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

ed42b5ac

sharding.py 13.6 KB