Merge pull request #171 from danielward27/equinox_errors

Test for updated equinox errors

pyproject.toml

@@ -16,7 +16,7 @@ license = { file = "LICENSE" }
 name = "flowjax"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "13.1.0"
+version = "13.1.1"
 
 [project.urls]
 repository = "https://github.com/danielward27/flowjax"

tests/test_bijections/test_bijection_utils.py

@@ -1,6 +1,6 @@
-import jax
 import jax.numpy as jnp
 import pytest
+from equinox import EquinoxRuntimeError
 
 from flowjax.bijections import Affine, Partial, Permute
 
@@ -30,5 +30,5 @@ def test_partial(idx, expected):
 
 
 def test_Permute_argcheck():
-    with pytest.raises(jax.lib.xla_extension.XlaRuntimeError):
+    with pytest.raises(EquinoxRuntimeError):
         Permute(jnp.array([0, 0]))

tests/test_wrappers.py

@@ -18,7 +18,7 @@
 
 def test_BijectionReparam():
 
-    with pytest.raises(jax.lib.xla_extension.XlaRuntimeError, match="Exp"):
+    with pytest.raises(eqx.EquinoxRuntimeError, match="Exp"):
         BijectionReparam(-jnp.ones(3), Exp())
 
     param = jnp.array([jnp.inf, 1, 2])