Implement JAX dispatch for IntegersRV

Author: ricardoV94

pytensor/link/jax/dispatch/random.py

@@ -179,6 +179,7 @@ def sample_fn(rng, size, dtype, *parameters):
 
 
 @jax_sample_fn.register(aer.RandIntRV)
+@jax_sample_fn.register(aer.IntegersRV)
 @jax_sample_fn.register(aer.UniformRV)
 def jax_sample_fn_uniform(op):
     """JAX implementation of random variables with uniform density.
@@ -188,6 +189,9 @@ def jax_sample_fn_uniform(op):
 
     """
     name = op.name
+    # IntegersRV is equivalent to RandintRV
+    if isinstance(op, aer.IntegersRV):
+        name = "randint"
     jax_op = getattr(jax.random, name)
 
     def sample_fn(rng, size, dtype, *parameters):

tests/link/jax/test_random.py

@@ -237,6 +237,22 @@ def test_random_updates(rng_ctor):
             "randint",
             lambda *args: args,
         ),
+        (
+            aer.integers,
+            [
+                set_test_value(
+                    at.lscalar(),
+                    np.array(0, dtype=np.int64),
+                ),
+                set_test_value(  # high-value necessary since test on cdf
+                    at.lscalar(),
+                    np.array(1000, dtype=np.int64),
+                ),
+            ],
+            (),
+            "randint",
+            lambda *args: args,
+        ),
         (
             aer.standard_normal,
             [],
@@ -376,7 +392,11 @@ def test_random_RandomVariable(rv_op, dist_params, base_size, cdf_name, params_c
         The parameters passed to the op.
 
     """
-    rng = shared(np.random.RandomState(29402))
+    if rv_op is aer.integers:
+        # Integers only accepts Generator, not RandomState
+        rng = shared(np.random.default_rng(29402))
+    else:
+        rng = shared(np.random.RandomState(29402))
     g = rv_op(*dist_params, size=(10_000,) + base_size, rng=rng)
     g_fn = function(dist_params, g, mode=jax_mode)
     samples = g_fn(