transformer_engine/jax/flax/module.py · 9f3e79bff824d3a9f10267dc414308011c87b093 · OpenDAS / TransformerEngine

"transformer_engine/jax/attention.py" did not exist on "b5a7c9f95e995236afae301366ec433c73b52690"

[JAX] Clamped Swiglu Integration (#2194) · b840898b

vthumbe1503 authored Oct 03, 2025


Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
*Jax integration for clamped swiglu. This is the continuation of PR which added Clamped Swiglu(used in GPT OSS) support in TE along with Pytorch integration. This PR hooks up the clamped swiglu and dswiglu's nvte APIs to TE Jax.

b840898b

module.py 50.2 KB

Replace module.py