transformer_engine/jax/cpp_extensions/activation.py · ca468ebe64c8af3fb49ef11cb47b09d07ff4fd2c · OpenDAS / TransformerEngine

[JAX] Use TE quant if TE fused act is disabled (#2374) · 389a6ba4

jberchtold-nvidia authored Nov 14, 2025



* Use TE quant if TE fused act is disabled
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Keep existing precision
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

389a6ba4

activation.py 56.5 KB

Replace activation.py