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
 
 
@@ -57,8 +58,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)
@@ -83,7 +83,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 RandomType 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)
@@ -130,12 +159,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
 
@@ -157,13 +184,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
 
@@ -175,12 +200,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
 
@@ -201,15 +224,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
 
@@ -226,13 +247,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
 
@@ -241,13 +260,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
 
@@ -256,17 +273,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
 
@@ -275,13 +290,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
 
@@ -290,13 +303,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
 
@@ -311,15 +322,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
 
@@ -334,15 +342,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
 
@@ -357,17 +362,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
 
 
@@ -14,6 +14,15 @@ def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
         from pytensor.link.jax.dispatch import jax_funcify
         from pytensor.tensor.random.type import RandomType
 
+        if any(
+            isinstance(inp.type, RandomType) and not isinstance(inp, SharedVariable)
+            for inp in fgraph.inputs
+        ):
+            warnings.warn(
+                "RandomTypes are implicitly converted to random PRNGKey arrays in JAX. "
+                "Input values should be provided in this format to avoid a conversion overhead."
+            )
+
         shared_rng_inputs = [
             inp
             for inp in fgraph.inputs
@@ -70,6 +79,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

@@ -5,11 +5,11 @@
 import pytensor
 import pytensor.tensor as at
 import pytensor.tensor.random as aer
-from pytensor.compile.function import function
 from pytensor.compile.sharedvalue import SharedVariable, shared
 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, random_state_type
 from pytensor.tensor.random.utils import RandomStream
 from tests.link.jax.test_basic import compare_jax_and_py, jax_mode, set_test_value
 
@@ -24,7 +24,40 @@ def random_function(*args, **kwargs):
     with pytest.warns(
         UserWarning, match=r"The RandomType SharedVariables \[.+\] will not be used"
     ):
-        return function(*args, **kwargs)
+        return pytensor.function(*args, **kwargs)
+
+
+@pytest.mark.parametrize("random_type", ("generator", "state"))
+def test_rng_io(random_type):
+    """Test explicit (non-shared) input and output RNG types in JAX."""
+    if random_type == "generator":
+        rng = random_generator_type("rng")
+        np_rng = np.random.default_rng(0)
+    else:
+        rng = random_state_type("rng")
+        np_rng = np.random.RandomState(0)
+    jx_rng = jax.random.PRNGKey(0)
+
+    next_rng, x = aer.normal(rng=rng).owner.outputs
+
+    with pytest.warns(
+        UserWarning,
+        match="RandomTypes are implicitly converted to random PRNGKey arrays",
+    ):
+        fn = pytensor.function([rng], [next_rng, x], mode="JAX")
+
+    # Inputs - RNG outputs
+    assert isinstance(fn(np_rng)[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(jx_rng)[1] == fn(jx_rng)[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(jx_rng)[0])[1] != fn(jx_rng)[1]
 
 
 def test_random_RandomStream():