fix: Use jax.errors API as jax._src removed from public namespace (…

…#2034) * Use jax.errors API over jax._src.errors as jax._src is removed from the jax public namespace in jax v0.3.18. - c.f. https://github.com/google/jax/releases/tag/jax-v0.3.18
scikit-hep · Sep 26, 2022 · d088b17 · d088b17
1 parent 2d791e8
commit d088b17
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_backends.py b/tests/test_backends.py
@@ -62,7 +62,7 @@ def example_op(x):
     with pytest.raises(
         (
             ValueError,
-            jax._src.errors.TracerArrayConversionError,
+            jax.errors.TracerArrayConversionError,
             jax.errors.ConcretizationTypeError,
         )
     ):
@@ -73,7 +73,7 @@ def example_op2(x):
         z = pyhf.tensorlib.sum(y)
         return z
 
-    with pytest.raises(jax._src.errors.TracerArrayConversionError):
+    with pytest.raises(jax.errors.TracerArrayConversionError):
         jax.jacrev(example_op2)(pyhf.tensorlib.astensor([2.0, 3.0]))