transformer_engine/jax/cpp_extensions/quantization.py · 3d0ea80a77d35eb07647c60aa3a3b54ce67c82c8 · OpenDAS / TransformerEngine

[JAX] `ScaledTensor1x` to store `amax` (#2117) · 3d0ea80a

Phuong Nguyen authored Aug 26, 2025



* added amax as an optional arg
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

3d0ea80a

quantization.py 34.6 KB

Replace quantization.py