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

[JAX] `dot_1_output` sharding constraint + use AXIS_IS_UNSHARDED (#2128) · a282136c

Phuong Nguyen authored Aug 27, 2025



* add dot_1_output sharding constraint + use AXIS_IS_UNSHARDED
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


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

a282136c

layernorm_mlp.py 16 KB

Replace layernorm_mlp.py