Typify RNG input variables in JAX linker

Author: ricardoV94

pytensor/link/basic.py

@@ -609,6 +609,9 @@ def create_thunk_inputs(self, storage_map: Dict[Variable, List[Any]]) -> List[An
     def jit_compile(self, fn: Callable) -> Callable:
         """JIT compile a converted ``FunctionGraph``."""
 
+    def typify(self, var: Variable):
+        return var
+
     def output_filter(self, var: Variable, out: Any) -> Any:
         """Apply a filter to the data output by a JITed function call."""
         return out
@@ -735,7 +738,7 @@ def make_all(self, input_storage=None, output_storage=None, storage_map=None):
         return (
             fn,
             [
-                Container(input, storage)
+                Container(self.typify(input), storage)
                 for input, storage in zip(fgraph.inputs, input_storage)
             ],
             [

pytensor/link/jax/dispatch/random.py

@@ -10,6 +10,7 @@
 import pytensor.tensor.random.basic as aer
 from pytensor.link.jax.dispatch.basic import jax_funcify, jax_typify
 from pytensor.link.jax.dispatch.shape import JAXShapeTuple
+from pytensor.tensor.random.type import RandomType
 from pytensor.tensor.shape import Shape, Shape_i
 
 
@@ -55,8 +56,7 @@ def jax_typify_RandomState(state, **kwargs):
     state = state.get_state(legacy=False)
     state["bit_generator"] = numpy_bit_gens[state["bit_generator"]]
     # XXX: Is this a reasonable approach?
-    state["jax_state"] = state["state"]["key"][0:2]
-    return state
+    return state["state"]["key"][0:2]
 
 
 @jax_typify.register(Generator)
@@ -81,7 +81,36 @@ def jax_typify_Generator(rng, **kwargs):
     state_32 = _coerce_to_uint32_array(state["state"]["state"])
     state["state"]["inc"] = inc_32[0] << 32 | inc_32[1]
     state["state"]["state"] = state_32[0] << 32 | state_32[1]
-    return state
+    return state["jax_state"]
+
+
+class RandomPRNGKeyType(RandomType[jax.random.PRNGKey]):
+    """JAX-compatible PRNGKey type.
+
+    This type is not exposed to users directly.
+
+    It is introduced by the JIT linker in place of any RandomStateType or RandomGeneratorType
+    input variables used in the original function. Nodes in the function graph will
+    still show the original types as inputs and outputs.
+    """
+
+    def filter(self, data, strict: bool = False, allow_downcast=None):
+        # PRNGs are just JAX Arrays, we assume this is a valid one!
+        if isinstance(data, jax.Array):
+            return data
+
+        if strict:
+            raise TypeError()
+
+        return jax_typify(data)
+
+
+random_prng_key_type = RandomPRNGKeyType()
+
+
+@jax_typify.register(RandomType)
+def jax_typify_RandomType(type):
+    return random_prng_key_type()
 
 
 @jax_funcify.register(aer.RandomVariable)
@@ -128,12 +157,10 @@ def jax_sample_fn_generic(op):
     name = op.name
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, *parameters):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         sample = jax_op(sampling_key, *parameters, shape=size, dtype=dtype)
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -155,13 +182,11 @@ def jax_sample_fn_loc_scale(op):
     name = op.name
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, *parameters):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         loc, scale = parameters
         sample = loc + jax_op(sampling_key, size, dtype) * scale
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -173,12 +198,10 @@ def jax_sample_fn_no_dtype(op):
     name = op.name
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, *parameters):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         sample = jax_op(sampling_key, *parameters, shape=size)
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -199,15 +222,13 @@ def jax_sample_fn_uniform(op):
         name = "randint"
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, *parameters):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         minval, maxval = parameters
         sample = jax_op(
             sampling_key, shape=size, dtype=dtype, minval=minval, maxval=maxval
         )
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -224,13 +245,11 @@ def jax_sample_fn_shape_rate(op):
     name = op.name
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, *parameters):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         (shape, rate) = parameters
         sample = jax_op(sampling_key, shape, size, dtype) / rate
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -239,13 +258,11 @@ def sample_fn(rng, size, dtype, *parameters):
 def jax_sample_fn_exponential(op):
     """JAX implementation of `ExponentialRV`."""
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, *parameters):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         (scale,) = parameters
         sample = jax.random.exponential(sampling_key, size, dtype) * scale
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -254,17 +271,15 @@ def sample_fn(rng, size, dtype, *parameters):
 def jax_sample_fn_t(op):
     """JAX implementation of `StudentTRV`."""
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, *parameters):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         (
             df,
             loc,
             scale,
         ) = parameters
         sample = loc + jax.random.t(sampling_key, df, size, dtype) * scale
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -273,13 +288,11 @@ def sample_fn(rng, size, dtype, *parameters):
 def jax_funcify_choice(op):
     """JAX implementation of `ChoiceRV`."""
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, *parameters):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         (a, p, replace) = parameters
         smpl_value = jax.random.choice(sampling_key, a, size, replace, p)
-        rng["jax_state"] = rng_key
-        return (rng, smpl_value)
+        return (rng_key, smpl_value)
 
     return sample_fn
 
@@ -288,13 +301,11 @@ def sample_fn(rng, size, dtype, *parameters):
 def jax_sample_fn_permutation(op):
     """JAX implementation of `PermutationRV`."""
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, *parameters):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
         (x,) = parameters
         sample = jax.random.permutation(sampling_key, x)
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -309,15 +320,12 @@ def jax_sample_fn_binomial(op):
 
     from numpyro.distributions.util import binomial
 
-    def sample_fn(rng, size, dtype, n, p):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, n, p):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
 
         sample = binomial(key=sampling_key, n=n, p=p, shape=size)
 
-        rng["jax_state"] = rng_key
-
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -332,15 +340,12 @@ def jax_sample_fn_multinomial(op):
 
     from numpyro.distributions.util import multinomial
 
-    def sample_fn(rng, size, dtype, n, p):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, n, p):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
 
         sample = multinomial(key=sampling_key, n=n, p=p, shape=size)
 
-        rng["jax_state"] = rng_key
-
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn
 
@@ -355,17 +360,14 @@ def jax_sample_fn_vonmises(op):
 
     from numpyro.distributions.util import von_mises_centered
 
-    def sample_fn(rng, size, dtype, mu, kappa):
-        rng_key = rng["jax_state"]
+    def sample_fn(rng_key, size, dtype, mu, kappa):
         rng_key, sampling_key = jax.random.split(rng_key, 2)
 
         sample = von_mises_centered(
             key=sampling_key, concentration=kappa, shape=size, dtype=dtype
         )
         sample = (sample + mu + np.pi) % (2.0 * np.pi) - np.pi
 
-        rng["jax_state"] = rng_key
-
-        return (rng, sample)
+        return (rng_key, sample)
 
     return sample_fn

pytensor/link/jax/linker.py

@@ -3,7 +3,7 @@
 from numpy.random import Generator, RandomState
 
 from pytensor.compile.sharedvalue import SharedVariable, shared
-from pytensor.graph.basic import Constant
+from pytensor.graph.basic import Constant, Variable
 from pytensor.link.basic import JITLinker
 
 
@@ -63,6 +63,11 @@ def jit_compile(self, fn):
         ]
         return jax.jit(fn, static_argnums=static_argnums)
 
+    def typify(self, var: Variable):
+        from pytensor.link.jax.dispatch import jax_typify
+
+        return jax_typify(var.type)
+
     def create_thunk_inputs(self, storage_map):
         from pytensor.link.jax.dispatch import jax_typify
 

tests/link/jax/test_random.py

@@ -12,6 +12,7 @@
 from pytensor.graph.basic import Constant
 from pytensor.graph.fg import FunctionGraph
 from pytensor.tensor.random.basic import RandomVariable
+from pytensor.tensor.random.type import random_generator_type
 from pytensor.tensor.random.utils import RandomStream
 from tests.link.jax.test_basic import compare_jax_and_py, jax_mode, set_test_value
 
@@ -22,6 +23,33 @@
 from pytensor.link.jax.dispatch.random import numpyro_available  # noqa: E402
 
 
+def test_rng_io():
+    rng = random_generator_type("rng")
+    next_rng, x = aer.normal(rng=rng).owner.outputs
+    fn = pytensor.function([rng], [next_rng, x], mode="JAX")
+
+    np_rng = np.random.default_rng(0)
+    np_rst = np.random.RandomState(1)
+    jx_rng = jax.random.PRNGKey(2)
+
+    # Inputs - RNG outputs
+    assert isinstance(fn(np_rng)[0], jax.Array)
+    assert isinstance(fn(np_rst)[0], jax.Array)
+    assert isinstance(fn(jx_rng)[0], jax.Array)
+
+    # Inputs - Value outputs
+    assert fn(np_rng)[1] == fn(np_rng)[1]
+    assert fn(np_rst)[1] == fn(np_rst)[1]
+    assert fn(jx_rng)[1] == fn(jx_rng)[1]
+    assert fn(np_rng)[1] != fn(np_rst)[1]
+    assert fn(np_rng)[1] != fn(jx_rng)[1]
+
+    # Chained Inputs - RNG / Value outputs
+    assert fn(fn(np_rng)[0])[1] != fn(np_rng)[1]
+    assert fn(fn(np_rst)[0])[1] != fn(np_rst)[1]
+    assert fn(fn(jx_rng)[0])[1] != fn(jx_rng)[1]
+
+
 def test_random_RandomStream():
     """Two successive calls of a compiled graph using `RandomStream` should
     return different values.