Revert unrelated code, remove deprecated code from a test

Author: jaanerik

tests/link/jax/test_scalar.py

@@ -38,13 +38,10 @@
 
 
 try:
-    import tensorflow_probability  # noqa: F401
-    from jax.interpreters.xla import (
-        pytype_aval_mappings,  # This is what's missing in new JAX  # noqa: F401
-    )
+    pass
 
     TFP_INSTALLED = True
-except (ModuleNotFoundError, AttributeError, ImportError):
+except ModuleNotFoundError:
     TFP_INSTALLED = False
 
 
@@ -163,7 +160,6 @@ def test_tfp_ops(op, test_values):
     compare_jax_and_py(inputs, [output], test_values)
 
 
-@pytest.mark.skipif(not TFP_INSTALLED, reason="Test requires tensorflow-probability")
 def test_betaincinv():
     a = vector("a", dtype="float64")
     b = vector("b", dtype="float64")
@@ -181,7 +177,6 @@ def test_betaincinv():
     )
 
 
-@pytest.mark.skipif(not TFP_INSTALLED, reason="Test requires tensorflow-probability")
 def test_gammaincinv():
     k = vector("k", dtype="float64")
     x = vector("x", dtype="float64")
@@ -190,7 +185,6 @@ def test_gammaincinv():
     compare_jax_and_py([k, x], [out], [np.array([5.5, 7.0]), np.array([0.25, 0.7])])
 
 
-@pytest.mark.skipif(not TFP_INSTALLED, reason="Test requires tensorflow-probability")
 def test_gammainccinv():
     k = vector("k", dtype="float64")
     x = vector("x", dtype="float64")

tests/link/jax/test_tensor_basic.py

@@ -226,8 +226,6 @@ def test_tri_nonconcrete():
 
     out = ptb.tri(m, n, k)
 
-    # The actual error the user will see should be jax.errors.ConcretizationTypeError, but
-    # the error handler raises an Attribute error first, so that's what this test needs to pass
     with pytest.raises(
         NotImplementedError,
         match=re.escape(

tests/tensor/rewriting/test_elemwise.py

@@ -1642,9 +1642,9 @@ def test_InplaceElemwiseOptimizer_bug():
     # with config.change_flags(tensor__insert_inplace_optimizer_validate_nb=10):
     rewrite_graph(fgraph, include=("inplace",))
 
-    with config.change_flags(tensor__insert_inplace_optimizer_validate_nb=1):
-        with pytest.warns(
-            FutureWarning,
-            match="tensor__insert_inplace_optimizer_validate_nb config is deprecated",
-        ):
-            rewrite_graph(fgraph, include=("inplace",))
+    pytensor.config.tensor__insert_inplace_optimizer_validate_nb = 1
+    with pytest.warns(
+        FutureWarning,
+        match="tensor__insert_inplace_optimizer_validate_nb config is deprecated",
+    ):
+        rewrite_graph(fgraph, include=("inplace",))

tests/tensor/test_basic.py

@@ -704,7 +704,7 @@ def test_masked_array_not_implemented(
 
 
 def check_alloc_runtime_broadcast(mode):
-    """Check we emit a clear error when runtime broadcasting would occur according to Numpy rules."""
+    """Check we emmit a clear error when runtime broadcasting would occur according to Numpy rules."""
     floatX = config.floatX
     x_v = vector("x", shape=(None,))