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

* Remove assertion for NO_MASK
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix JAX distributed unit tests name
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>

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_*
 

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

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,