Default to JAX test mode in random tests

Author: ricardoV94

tests/link/jax/test_random.py

@@ -27,7 +27,7 @@
 from pytensor.link.jax.dispatch.random import numpyro_available  # noqa: E402
 
 
-def compile_random_function(*args, mode="JAX", **kwargs):
+def compile_random_function(*args, mode=jax_mode, **kwargs):
     with pytest.warns(
         UserWarning, match=r"The RandomType SharedVariables \[.+\] will not be used"
     ):
@@ -42,7 +42,7 @@ def test_random_RandomStream():
     srng = RandomStream(seed=123)
     out = srng.normal() - srng.normal()
 
-    fn = compile_random_function([], out, mode=jax_mode)
+    fn = compile_random_function([], out)
     jax_res_1 = fn()
     jax_res_2 = fn()
 
@@ -55,7 +55,7 @@ def test_random_updates(rng_ctor):
     rng = shared(original_value, name="original_rng", borrow=False)
     next_rng, x = pt.random.normal(name="x", rng=rng).owner.outputs
 
-    f = compile_random_function([], [x], updates={rng: next_rng}, mode=jax_mode)
+    f = compile_random_function([], [x], updates={rng: next_rng})
     assert f() != f()
 
     # Check that original rng variable content was not overwritten when calling jax_typify
@@ -475,7 +475,7 @@ def test_random_RandomVariable(rv_op, dist_params, base_size, cdf_name, params_c
     """
     rng = shared(np.random.default_rng(29403))
     g = rv_op(*dist_params, size=(10000, *base_size), rng=rng)
-    g_fn = compile_random_function(dist_params, g, mode=jax_mode)
+    g_fn = compile_random_function(dist_params, g)
     samples = g_fn(
         *[
             i.tag.test_value
@@ -517,7 +517,7 @@ def test_size_implied_by_broadcasted_parameters(rv_fn):
     param_that_implies_size = pt.matrix("param_that_implies_size", shape=(None, None))
 
     rv = rv_fn(param_that_implies_size)
-    draws = rv.eval({param_that_implies_size: np.zeros((2, 2))}, mode=jax_mode)
+    draws = rv.eval({param_that_implies_size: np.zeros((2, 2))})
 
     assert draws.shape == (2, 2)
     assert np.unique(draws).size == 4
@@ -527,7 +527,7 @@ def test_size_implied_by_broadcasted_parameters(rv_fn):
 def test_random_bernoulli(size):
     rng = shared(np.random.default_rng(123))
     g = pt.random.bernoulli(0.5, size=(1000, *size), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     np.testing.assert_allclose(samples.mean(axis=0), 0.5, 1)
 
@@ -538,7 +538,7 @@ def test_random_mvnormal():
     mu = np.ones(4)
     cov = np.eye(4)
     g = pt.random.multivariate_normal(mu, cov, size=(10000,), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     np.testing.assert_allclose(samples.mean(axis=0), mu, atol=0.1)
 
@@ -553,7 +553,7 @@ def test_random_mvnormal():
 def test_random_dirichlet(parameter, size):
     rng = shared(np.random.default_rng(123))
     g = pt.random.dirichlet(parameter, size=(1000, *size), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     np.testing.assert_allclose(samples.mean(axis=0), 0.5, 1)
 
@@ -562,7 +562,7 @@ def test_random_choice():
     # `replace=True` and `p is None`
     rng = shared(np.random.default_rng(123))
     g = pt.random.choice(np.arange(4), size=10_000, rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     assert samples.shape == (10_000,)
     # Elements are picked at equal frequency
@@ -571,7 +571,7 @@ def test_random_choice():
     # `replace=True` and `p is not None`
     rng = shared(np.random.default_rng(123))
     g = pt.random.choice(4, p=np.array([0.0, 0.5, 0.0, 0.5]), size=(5, 2), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     assert samples.shape == (5, 2)
     # Only odd numbers are picked
@@ -580,7 +580,7 @@ def test_random_choice():
     # `replace=False` and `p is None`
     rng = shared(np.random.default_rng(123))
     g = pt.random.choice(np.arange(100), replace=False, size=(2, 49), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     assert samples.shape == (2, 49)
     # Elements are unique
@@ -595,7 +595,7 @@ def test_random_choice():
         rng=rng,
         replace=False,
     )
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     assert samples.shape == (3,)
     # Elements are unique
@@ -607,14 +607,14 @@ def test_random_choice():
 def test_random_categorical():
     rng = shared(np.random.default_rng(123))
     g = pt.random.categorical(0.25 * np.ones(4), size=(10000, 4), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     assert samples.shape == (10000, 4)
     np.testing.assert_allclose(samples.mean(axis=0), 6 / 4, 1)
 
     # Test zero probabilities
     g = pt.random.categorical([0, 0.5, 0, 0.5], size=(1000,), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     assert samples.shape == (1000,)
     assert np.all(samples % 2 == 1)
@@ -624,7 +624,7 @@ def test_random_permutation():
     array = np.arange(4)
     rng = shared(np.random.default_rng(123))
     g = pt.random.permutation(array, rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     permuted = g_fn()
     with pytest.raises(AssertionError):
         np.testing.assert_allclose(array, permuted)
@@ -647,7 +647,7 @@ def test_random_geometric():
     rng = shared(np.random.default_rng(123))
     p = np.array([0.3, 0.7])
     g = pt.random.geometric(p, size=(10_000, 2), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     np.testing.assert_allclose(samples.mean(axis=0), 1 / p, rtol=0.1)
     np.testing.assert_allclose(samples.std(axis=0), np.sqrt((1 - p) / p**2), rtol=0.1)
@@ -658,7 +658,7 @@ def test_negative_binomial():
     n = np.array([10, 40])
     p = np.array([0.3, 0.7])
     g = pt.random.negative_binomial(n, p, size=(10_000, 2), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     np.testing.assert_allclose(samples.mean(axis=0), n * (1 - p) / p, rtol=0.1)
     np.testing.assert_allclose(
@@ -672,7 +672,7 @@ def test_binomial():
     n = np.array([10, 40])
     p = np.array([0.3, 0.7])
     g = pt.random.binomial(n, p, size=(10_000, 2), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     np.testing.assert_allclose(samples.mean(axis=0), n * p, rtol=0.1)
     np.testing.assert_allclose(samples.std(axis=0), np.sqrt(n * p * (1 - p)), rtol=0.1)
@@ -687,7 +687,7 @@ def test_beta_binomial():
     a = np.array([1.5, 13])
     b = np.array([0.5, 9])
     g = pt.random.betabinom(n, a, b, size=(10_000, 2), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     np.testing.assert_allclose(samples.mean(axis=0), n * a / (a + b), rtol=0.1)
     np.testing.assert_allclose(
@@ -721,7 +721,7 @@ def test_vonmises_mu_outside_circle():
     mu = np.array([-30, 40])
     kappa = np.array([100, 10])
     g = pt.random.vonmises(mu, kappa, size=(10_000, 2), rng=rng)
-    g_fn = compile_random_function([], g, mode=jax_mode)
+    g_fn = compile_random_function([], g)
     samples = g_fn()
     np.testing.assert_allclose(
         samples.mean(axis=0), (mu + np.pi) % (2.0 * np.pi) - np.pi, rtol=0.1
@@ -819,15 +819,15 @@ def test_random_concrete_shape():
     rng = shared(np.random.default_rng(123))
     x_pt = pt.dmatrix()
     out = pt.random.normal(0, 1, size=x_pt.shape, rng=rng)
-    jax_fn = compile_random_function([x_pt], out, mode=jax_mode)
+    jax_fn = compile_random_function([x_pt], out)
     assert jax_fn(np.ones((2, 3))).shape == (2, 3)
 
 
 def test_random_concrete_shape_from_param():
     rng = shared(np.random.default_rng(123))
     x_pt = pt.dmatrix()
     out = pt.random.normal(x_pt, 1, rng=rng)
-    jax_fn = compile_random_function([x_pt], out, mode=jax_mode)
+    jax_fn = compile_random_function([x_pt], out)
     assert jax_fn(np.ones((2, 3))).shape == (2, 3)
 
 
@@ -846,7 +846,7 @@ def test_random_concrete_shape_subtensor():
     rng = shared(np.random.default_rng(123))
     x_pt = pt.dmatrix()
     out = pt.random.normal(0, 1, size=x_pt.shape[1], rng=rng)
-    jax_fn = compile_random_function([x_pt], out, mode=jax_mode)
+    jax_fn = compile_random_function([x_pt], out)
     assert jax_fn(np.ones((2, 3))).shape == (3,)
 
 
@@ -862,7 +862,7 @@ def test_random_concrete_shape_subtensor_tuple():
     rng = shared(np.random.default_rng(123))
     x_pt = pt.dmatrix()
     out = pt.random.normal(0, 1, size=(x_pt.shape[0],), rng=rng)
-    jax_fn = compile_random_function([x_pt], out, mode=jax_mode)
+    jax_fn = compile_random_function([x_pt], out)
     assert jax_fn(np.ones((2, 3))).shape == (2,)
 
 
@@ -873,7 +873,7 @@ def test_random_concrete_shape_graph_input():
     rng = shared(np.random.default_rng(123))
     size_pt = pt.scalar()
     out = pt.random.normal(0, 1, size=size_pt, rng=rng)
-    jax_fn = compile_random_function([size_pt], out, mode=jax_mode)
+    jax_fn = compile_random_function([size_pt], out)
     assert jax_fn(10).shape == (10,)