Skip tri test in latest version of JAX

Related to jax-ml/jax#22751
pymc-devs · Aug 4, 2024 · de06ff6 · de06ff6
1 parent f9bf218
commit de06ff6
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tests/link/jax/test_tensor_basic.py b/tests/link/jax/test_tensor_basic.py
@@ -218,6 +218,10 @@ def test_tri():
     compare_jax_and_py(fgraph, [])
 
 
+@pytest.mark.skipif(
+    jax.__version__ == "0.4.31",
+    reason="https://github.com/google/jax/issues/22751",
+)
 def test_tri_nonconcrete():
     """JAX cannot JIT-compile `jax.numpy.tri` when arguments are not concrete values."""