Update equinox to >=0.11.3.

This was to fix breaking changes in recent versions of JAX (0.4.32 and 0.4.33) which are supported.

PiperOrigin-RevId: 682333073

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "numpy>=1.24.1",
     "setuptools;python_version>='3.10'",
     "chex>=0.1.85",
-    "equinox @ git+https://github.com/patrick-kidger/equinox@1e601672d38d2c4d483535070a3572d8e8508a20",
+    "equinox>=0.11.3",
     "PyYAML>=6.0.1",
     "xarray>=2023.12.0",
     "netcdf4>=1.6.5,<1.7.1",

torax/config/tests/runtime_params_slice.py

@@ -305,7 +305,7 @@ def test_wext_in_dynamic_runtime_params_cannot_be_negative(self):
     )
     np.testing.assert_allclose(jext.wext, 0.0)
     # But negative values will cause an error.
-    with self.assertRaises(jax.lib.xla_client.XlaRuntimeError):
+    with self.assertRaises(RuntimeError):
       dcs_provider(t=1.0,)
 
   @parameterized.parameters(

torax/sources/tests/external_current_source.py

@@ -15,7 +15,6 @@
 """Tests for external_current_source."""
 
 from absl.testing import absltest
-import jax
 from torax import geometry
 from torax.config import runtime_params as general_runtime_params
 from torax.config import runtime_params_slice
@@ -79,7 +78,7 @@ def test_invalid_source_types_raise_errors(self):
     source = source_builder()
     for unsupported_mode in self._unsupported_modes:
       with self.subTest(unsupported_mode.name):
-        with self.assertRaises(jax.lib.xla_client.XlaRuntimeError):
+        with self.assertRaises(RuntimeError):
           source_builder.runtime_params.mode = unsupported_mode
           dynamic_slice = runtime_params_slice.DynamicRuntimeParamsSliceProvider(
               runtime_params,

torax/sources/tests/qei_source.py

@@ -16,7 +16,6 @@
 
 import dataclasses
 from absl.testing import absltest
-import jax
 from torax import core_profile_setters
 from torax import geometry
 from torax.config import runtime_params as general_runtime_params
@@ -106,7 +105,7 @@ def test_invalid_source_types_raise_errors(self):
     )
     for unsupported_mode in self._unsupported_modes:
       with self.subTest(unsupported_mode.name):
-        with self.assertRaises(jax.lib.xla_client.XlaRuntimeError):
+        with self.assertRaises(RuntimeError):
           dynamic_slice = runtime_params_slice.DynamicRuntimeParamsSliceProvider(
               runtime_params,
               sources={

torax/sources/tests/source.py

@@ -17,7 +17,6 @@
 import dataclasses
 from absl.testing import absltest
 from absl.testing import parameterized
-import jax
 from jax import numpy as jnp
 import numpy as np
 from torax import core_profile_setters
@@ -195,7 +194,7 @@ def test_unsupported_modes_raise_errors(self):
         source_models=source_models,
     )
     # But calling requesting ZERO shouldn't work.
-    with self.assertRaises(jax.lib.xla_client.XlaRuntimeError):
+    with self.assertRaises(RuntimeError):
       source.get_value(
           dynamic_runtime_params_slice=dynamic_runtime_params_slice,
           dynamic_source_runtime_params=dynamic_runtime_params_slice.sources[

torax/sources/tests/test_lib.py

@@ -158,7 +158,7 @@ def test_invalid_source_types_raise_errors(self):
           )
       )
       with self.subTest(unsupported_mode.name):
-        with self.assertRaises(jax.lib.xla_client.XlaRuntimeError):
+        with self.assertRaises(RuntimeError):
           source.get_value(
               dynamic_runtime_params_slice=dynamic_runtime_params_slice,
               dynamic_source_runtime_params=dynamic_runtime_params_slice.sources[
@@ -248,7 +248,7 @@ def test_invalid_source_types_raise_errors(self):
           )
       )
       with self.subTest(unsupported_mode.name):
-        with self.assertRaises(jax.lib.xla_client.XlaRuntimeError):
+        with self.assertRaises(RuntimeError):
           source.get_value(
               dynamic_runtime_params_slice=dynamic_runtime_params_slice,
               dynamic_source_runtime_params=dynamic_runtime_params_slice.sources[

torax/tests/jax_utils.py

@@ -16,7 +16,6 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import jax
 from jax import numpy as jnp
 from torax import jax_utils
 
@@ -29,8 +28,8 @@ def _should_error(self):
     x = jnp.array(0)
     cond = x == 0
 
-    with self.assertRaises(jax.lib.xla_extension.XlaRuntimeError):
-      x = jax_utils.error_if(x, cond, msg="")
+    with self.assertRaises(RuntimeError):
+      jax_utils.error_if(x, cond, msg="")
 
   def _should_not_error(self):
     """Call error_if, expecting it to be disabled.
@@ -41,7 +40,7 @@ def _should_not_error(self):
     x = jnp.array(0)
     cond = x == 0
 
-    x = jax_utils.error_if(x, cond, msg="")
+    jax_utils.error_if(x, cond, msg="")
 
   def test_enable_errors(self):
     """Test that jax_utils.enable_errors enables / disables errors."""