[JAX] Fix JAX distributed unit tests (#521)

* Remove assertion for NO_MASK Signed-off-by: Reese Wang <[email protected]> * Fix JAX distributed unit tests name Signed-off-by: Reese Wang <[email protected]> --------- Signed-off-by: Reese Wang <[email protected]>
NVIDIA · Nov 20, 2023 · ea43b18 · ea43b18
1 parent 6159af4
commit ea43b18
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 5 deletions.
diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh
@@ -5,5 +5,5 @@
 set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
-pytest -Wignore -v $TE_PATH/tests/jax/test_distributed_custom_ops.py
+pytest -Wignore -v $TE_PATH/tests/jax/test_distributed_*
 
diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh
@@ -5,7 +5,7 @@
 set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
-pytest -Wignore -v $TE_PATH/tests/jax
+pytest -Wignore -v $TE_PATH/tests/jax -k 'not distributed'
 
 pip install -r $TE_PATH/examples/jax/mnist/requirements.txt
 pip install -r $TE_PATH/examples/jax/encoder/requirements.txt

diff --git a/transformer_engine/jax/fused_attn.py b/transformer_engine/jax/fused_attn.py
@@ -54,9 +54,6 @@ def self_fused_attn(qkv: jnp.ndarray, bias: jnp.ndarray, mask: jnp.ndarray, seed
     """
     Self fused attention wrapper
     """
-    assert attn_mask_type is not AttnMaskType.NO_MASK, \
-        "Currently not support AttnMaskType.NO_MASK."
-
     output = _self_fused_attn(qkv,
                               bias,
                               mask,