examples/jax/mnist/test_single_gpu_mnist.py · 962d9c53423e604f3c22c3ad634bc5a0d66e4f7c · 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_single_gpu_mnist.py 11.4 KB

Replace test_single_gpu_mnist.py