transformer_engine/jax/quantize/quantizer.py · ef28c86582b5539655e7286366c699bec99284bc · OpenDAS / TransformerEngine

[JAX] Make SR rng state always 2D (num_devices, 4) to fix partitioning issue (#2294) · e2f2a0b4

jberchtold-nvidia authored Oct 23, 2025



* Make SR rng state always 2D (num_devices, 4)
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* fix pure-jax impl
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* fix test shape
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

---------
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

e2f2a0b4

quantizer.py 46.2 KB

Replace quantizer.py