Fixed some unit tests that failed on jax0.6.2 python 3.12

[eshen1991]

Fixed some unit tests that failed on jax0.6.2 python 3.12(jax0.6.2 requires python >=3.12) and validated that the fix also works for jax0.5.3(python3.10). The failures are due to AttributeErrors that are caused by removal in jax0.6.2 or python 3.12 of deprecated functions/properties that are still available in jax0.5.3/python3.10

Fixed Issues

1. AttributeError: module 'jax._src.ad_checkpoint' has no attribute '_optimization_barrier'

   • use lax.lax.opitmization_barrier instead of jax._src._optimizer_barrier, which is private property and is removed in 0.6.2

2. AttributeError: 'TestUtils' object has no attribute 'assertDictContainsSubset'. Did you mean: 'assertContainsSubset' assertDictContainsSubset is removed form unittest lib in python3.12.

   • Replace it with assertTrue --- self.assertDictContainsSubset({"text": True}, mock_popen.call_args[1]) +++ self.assertTrue({"text": True}.items() <= mock_popen.call_args[1].items())

3. AttributeError: module 'jaxlib' has no attribute 'xla_extension'

   • change log:
     jax.lib.xla_client, the previously-deprecated Device and XlaRuntimeError symbols have been removed; instead use jax.Device and jax.errors.JaxRuntimeError respectively.

Testing

Unit test results in axlearn-arc

[eshen1991]

pre-commit run -a worked locally

[eshen1991]

• Running pytype -j auto axlearn locally succeeded.

(axlearn-310) ➜  axlearn git:(eshen/jax-0.5.3) ✗ pytype -j auto axlearn
Computing dependencies
2025-08-29 18:37:28.019778: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2025-08-29 18:37:28.042197: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2025-08-29 18:37:28.049003: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-08-29 18:37:29.384596: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
WARNING fairseq is not installed -- param conversion for fairseq modules will fail.
jax version=0.5.3
Analyzing 595 sources with 0 local dependencies
ninja: Entering directory `.pytype'
[600/600] check axlearn.common.convolution_test
Leaving directory '.pytype'
Success: no errors found

• Running locally pytype -j auto . as pre-commit check does failed with ninja error, which is different from the pre-commit check error: Process completed with exit code 1.

(axlearn-310) ➜  axlearn git:(eshen/jax-0.5.3) ✗ pytype -j auto .      
Computing dependencies
2025-08-29 19:33:55.897094: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2025-08-29 19:33:55.920832: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2025-08-29 19:33:55.929141: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-08-29 19:33:57.733564: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
WARNING fairseq is not installed -- param conversion for fairseq modules will fail.
jax version=0.5.3
Analyzing 1191 sources with 0 local dependencies
ninja: Entering directory `.pytype'
ninja: error: build.ninja:757: multiple rules generate /usr/local/google/home/ericshen/git/axlearn/.pytype/pyi/axlearn/experiments/audio/conformer/librispeech_trainer.pyi

Leaving directory '.pytype'