[JAX] Add L2_jax_distributed_unittest (#2060)

* Add L2_jax_distributed_unittest
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Add L1 entry for NORM_INPUT_SHAPES that was missing
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

---------
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

qa/L1_jax_distributed_unittest/test.sh

@@ -8,4 +8,4 @@ set -xe
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"
 
-python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*

qa/L2_jax_distributed_unittest/test.sh

+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+set -xe
+
+: ${TE_PATH:=/opt/transformerengine}
+: ${XML_LOG_DIR:=/logs}
+mkdir -p "$XML_LOG_DIR"
+
+NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*

tests/jax/test_distributed_layernorm.py

@@ -25,6 +25,7 @@ DTYPES = [jnp.bfloat16, jnp.float32]
 
 NORM_INPUT_SHAPES = {
     "L0": [[64, 64]],
+    "L1": [[64, 64]],
     "L2": [[64, 64]],
 }