[JAX] Enable TE/JAX test timings in CI (#2475)

enable jax test timings in qa scripts
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

qa/L0_jax_distributed_unittest/test.sh

@@ -16,6 +16,8 @@ function test_fail() {
 RET=0
 FAILED_CASES=""
 
+export NVTE_JAX_TEST_TIMING=1
+
 : ${TE_PATH:=/opt/transformerengine}
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"

qa/L0_jax_unittest/test.sh

@@ -18,6 +18,8 @@ function test_fail() {
 RET=0
 FAILED_CASES=""
 
+export NVTE_JAX_TEST_TIMING=1
+
 pip3 install "nltk>=3.8.2" || error_exit "Failed to install nltk"
 pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
 

qa/L1_jax_distributed_unittest/test.sh

@@ -11,6 +11,8 @@ function test_fail() {
 RET=0
 FAILED_CASES=""
 
+export NVTE_JAX_TEST_TIMING=1
+
 : ${TE_PATH:=/opt/transformerengine}
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"

qa/L2_jax_distributed_unittest/test.sh

@@ -4,6 +4,8 @@
 
 set -xe
 
+export NVTE_JAX_TEST_TIMING=1
+
 : ${TE_PATH:=/opt/transformerengine}
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"

qa/L2_jax_unittest/test.sh

@@ -18,6 +18,8 @@ function test_fail() {
 RET=0
 FAILED_CASES=""
 
+export NVTE_JAX_TEST_TIMING=1
+
 pip3 install "nltk>=3.8.2" || error_exit "Failed to install nltk"
 pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"