transformer_engine/jax/csrc/extensions/activation.cpp · 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

activation.cpp 26.5 KB

Replace activation.cpp