tests/jax/test_custom_call_compute.py · ab3e5a92a1199e39b95097253dbf1ea7ef3b3a83 · OpenDAS / TransformerEngine

[JAX] Scaling Enum Abstracting (#1655) · 962d9c53

Phuong Nguyen authored Apr 09, 2025



* scaling enum abstract

* rm NVTE_ from ScalingMode names

* rework scaling mode enum in grouped gemm

* fix norm sharding

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

962d9c53

test_custom_call_compute.py 49.9 KB

Replace test_custom_call_compute.py