transformer_engine/jax/cpp_extensions/quantization.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

quantization.py 38 KB

Replace quantization.py