transformer_engine/jax/layernorm_mlp.py · 127b6d3ab3088c403f3b38b8405b70fc33ee3f34 · OpenDAS / TransformerEngine

[JAX] Activation/Normalization to output amax for later quantization in CurrentScaling (#2238) · 127b6d3a

Phuong Nguyen authored Oct 07, 2025



* reuse amax for current scaling
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

127b6d3a

layernorm_mlp.py 19.3 KB

Replace layernorm_mlp.py