tests/jax/test_distributed_layernorm_mlp.py · cae1c436027cc028ee49c83463141ba67f0adca0 · OpenDAS / TransformerEngine

[JAX] TE Gemm custom call clean up (#2030) · cae1c436

Phuong Nguyen authored Aug 07, 2025



* rm batch_dim, sequence_dim, sequence_parallel_output
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* rm lhs_quantized_colwise and rhs_quantized_colwise
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* rm unnecessary transpose_batch_sequence arg from some modules
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

cae1c436

test_distributed_layernorm_mlp.py 18.7 KB

Replace test_distributed_layernorm_mlp.py