transformer_engine/jax/quantize/scaling_modes.py · 3f875fb57fcf2872d238f8c7cb199b171c424536 · OpenDAS / TransformerEngine

[JAX] Restore Shardy Rule with CompoundFactor (#2167) · 3f875fb5

Phuong Nguyen authored Sep 23, 2025



* Rework shardy rules

* WAR for compound factor=1
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

3f875fb5

scaling_modes.py 29.5 KB

Replace scaling_modes.py