transformer_engine/jax/layernorm_mlp.py · dd083bdfca296294322f2b5f8143ce54ee15db22 · OpenDAS / TransformerEngine

Phuong Nguyen authored Aug 06, 2025



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

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

ed42b5ac

layernorm_mlp.py 17.5 KB