Handle shared variables that diverge across backends

doc/tutorial/prng.rst

@@ -251,9 +251,12 @@ Inplace optimization
 
 As mentioned, RandomVariable Ops default to making a copy of the input RNG before using it, which can be quite slow.
 
+We compile with the Python linker here so the dprint shows the plain RandomVariable form. Some backends specialize it (e.g. Numba wraps it in a ``RandomVariableWithCoreShape``, see the Numba section below).
+
+>>> py_mode = pytensor.Mode(linker="py", optimizer="fast_run")
 >>> rng = pt.random.shared_rng(np.random.default_rng(123), name="rng")
 >>> next_rng, x = pt.random.uniform(rng=rng, return_next_rng=True)
->>> f = pytensor.function([], x)
+>>> f = pytensor.function([], x, mode=py_mode)
 >>> _ = pytensor.dprint(f, print_destroy_map=True)
 uniform_rv{"(),()->()"}.1 [id A] 0
  ├─ rng [id B]
@@ -291,7 +294,7 @@ Users should never do this directly!
 >>> x.owner.op.inplace
 True
 
->>> inplace_f = pytensor.function([], x, accept_inplace=True)
+>>> inplace_f = pytensor.function([], x, accept_inplace=True, mode=py_mode)
 >>> _ = pytensor.dprint(inplace_f, print_destroy_map=True)
 uniform_rv{"(),()->()"}.1 [id A] d={0: [0]} 0
  ├─ rng [id B]
@@ -319,7 +322,7 @@ The first case is true when a user uses the `mutable` `kwarg` directly.
 >>> rng = pt.random.rng("rng")
 >>> next_rng, x = rng.uniform()
 >>> with pytensor.config.change_flags(optimizer_verbose=True): # doctest: +ELLIPSIS
-...     inplace_f = pytensor.function([In(rng, mutable=True)], [x])
+...     inplace_f = pytensor.function([In(rng, mutable=True)], [x], mode=py_mode)
 rewriting: rewrite random_make_inplace replaces ...
 >>> _ = inplace_f.dprint(print_destroy_map=True)
 uniform_rv{"(),()->()"}.1 [id A] d={0: [0]} 0
@@ -334,7 +337,7 @@ In this case, a RandomVariable is allowed to modify the RNG because the shared v
 >>> rng = pt.random.shared_rng(np.random.default_rng(123), name="rng")
 >>> next_rng, x = rng.uniform()
 >>>
->>> inplace_f = pytensor.function([], [x], updates={rng: next_rng})
+>>> inplace_f = pytensor.function([], [x], updates={rng: next_rng}, mode=py_mode)
 >>> _ = inplace_f.dprint(print_destroy_map=True)
 uniform_rv{"(),()->()"}.1 [id A] d={0: [0]} 0
  ├─ rng [id B]
@@ -529,19 +532,19 @@ NumPy random generators can be natively used with the Numba backend.
  ├─ [] [id B] <Vector(int64, shape=(0,))>
  ├─ randomstate_rng [id C] <RandomGeneratorType>
  ├─ NoneConst{None} [id D] <NoneTypeT>
- ├─ 0.0 [id E] <Scalar(float32, shape=())>
- └─ 1.0 [id F] <Scalar(float32, shape=())>
+ ├─ 0.0 [id E] <Scalar(float64, shape=())>
+ └─ 1.0 [id F] <Scalar(float64, shape=())>
 [normal_rv{"(),()->()"}].0 [id A] <RandomGeneratorType> 0
  └─ ···
 <BLANKLINE>
 Inner graphs:
 <BLANKLINE>
 [normal_rv{"(),()->()"}] [id A]
  ← normal_rv{"(),()->()"}.0 [id G] <RandomGeneratorType>
-    ├─ *1-<RandomGeneratorType> [id H] <RandomGeneratorType>
-    ├─ *2-<NoneTypeT> [id I] <NoneTypeT>
-    ├─ *3-<Scalar(float32, shape=())> [id J] <Scalar(float32, shape=())>
-    └─ *4-<Scalar(float32, shape=())> [id K] <Scalar(float32, shape=())>
+    ├─ i1 [id H] <RandomGeneratorType>
+    ├─ i2 [id I] <NoneTypeT>
+    ├─ i3 [id J] <Scalar(float64, shape=())>
+    └─ i4 [id K] <Scalar(float64, shape=())>
  ← normal_rv{"(),()->()"}.1 [id G] <Scalar(float64, shape=())>
     └─ ···
 
@@ -553,51 +556,63 @@ Numba converts regular RandomVariable to RandomVariableWithCoreShape Ops, subtly
 >>> type(numba_fn.maker.fgraph.outputs[0].owner.op)
 <class 'pytensor.tensor.random.op.RandomVariableWithCoreShape'>
 
-JAX (and MLX)
--------------
-
-JAX (and MLX) use a different type of PRNG than those of NumPy. This means that the standard shared RNGs cannot be used directly in graphs transpiled to those backends.
+JAX
+---
 
-Instead a copy of the Shared RNG variable is made, and its bit generator state is expanded with a jax_state/mlx_state entry. This is what's actually used by the random variables in those backends.
+JAX uses a different type of PRNG than those of NumPy: a NumPy ``Generator`` holds a PCG64 bit generator state, while JAX advances a threefry key.
+Because a single PyTensor graph can be transpiled to multiple backends, the shared RNG cannot store one representation that all backends use natively.
 
-In general, update rules are still respected, but they won't update/rely on the original shared variable.
+Instead, PyTensor keeps a per-backend *view* of the shared RNG's state (the JAX view expands the bit generator state with a ``jax_state`` entry).
+The shared variable is used directly in the graph, and reads and writes are reconciled so that ``get_value`` always observes the latest advancement and ``set_value`` and ``updates`` always take effect, regardless of which backend last touched the RNG.
+Reconciliation keeps the bookkeeping consistent, but crossing backends still reseeds the stream (see the warning below).
 
->>> import jax
 >>> rng = pt.random.shared_rng(np.random.default_rng(123), name="rng")
 >>> next_rng, x = rng.uniform()
 >>> jax_fn = pytensor.function([], [x], updates={rng: next_rng}, mode="JAX")
 >>> _ = pytensor.dprint(jax_fn, print_type=True) # doctest: +ELLIPSIS
 uniform_rv{"(),()->()"}.1 [id A] <Scalar(float64, shape=())> 0
- ├─ RNG(<Generator(PCG64) at 0x...>) [id B] <RandomGeneratorType>
+ ├─ rng [id B] <RandomGeneratorType>
  ├─ NoneConst{None} [id C] <NoneTypeT>
  ├─ 0.0 [id D] <Scalar(float32, shape=())>
  └─ 1.0 [id E] <Scalar(float32, shape=())>
 uniform_rv{"(),()->()"}.0 [id A] <RandomGeneratorType> 0
  └─ ···
 
 >>> print(jax_fn(), jax_fn())
-[Array(0.89545652, dtype=float64)] [Array(0.38045988, dtype=float64)]
+[Array(0.91459085, dtype=float64)] [Array(0.01298661, dtype=float64)]
 
->>> rng.set_value(np.random.default_rng(123))  # No effect on the jax evaluation
+``set_value`` propagates to the JAX evaluation, so reseeding the shared RNG resets the stream:
+
+>>> rng.set_value(np.random.default_rng(123))
 >>> print(jax_fn(), jax_fn())
-[Array(0.98049127, dtype=float64)] [Array(0.39260106, dtype=float64)]
-
->>> [jax_rng] = jax_fn.input_storage[0].storage
->>> jax_rng  # doctest: +NORMALIZE_WHITESPACE
-{'bit_generator': Array(1, dtype=int64, weak_type=True),
- 'has_uint32': Array(0, dtype=int64, weak_type=True),
- 'jax_state': Array([4091271363, 1319784711], dtype=uint32),
- 'state': {'inc': Array(651939783, dtype=uint32),
-  'state': Array(1542324465, dtype=uint32)},
- 'uinteger': Array(0, dtype=int64, weak_type=True)}
-
->>> [jax_rng] = jax_fn.input_storage[0].storage
->>> jax_rng["jax_state"] = jax.random.PRNGKey(0)
+[Array(0.91459085, dtype=float64)] [Array(0.01298661, dtype=float64)]
+
+.. warning::
+
+    Because the algorithms differ, once one backend has drawn from the RNG, another backend cannot continue producing that same sequence of numbers.
+    It can only reseed its own generator from the current state, which is deterministic but starts a new, unrelated stream.
+    This reseed loses the stream in either direction, so crossing backends emits a ``UserWarning``. Within a single backend nothing is converted and the RNG advances at zero cost.
+
+If a function's backend-specific updates are private (it only needs to observe host ``set_value`` and never has to feed its advanced state back to other backends or ``get_value``), pass ``readback=False`` to :func:`pytensor.function`.
+This drops all per-call cross-backend bookkeeping from the compiled function.
+
+>>> rng = pt.random.shared_rng(np.random.default_rng(123), name="rng")
+>>> next_rng, x = rng.uniform()
+>>> jax_fn = pytensor.function([], [x], updates={rng: next_rng}, mode="JAX", readback=False)
 >>> print(jax_fn(), jax_fn())
-[Array(0.08062437, dtype=float64)] [Array(0.67119299, dtype=float64)]
+[Array(0.91459085, dtype=float64)] [Array(0.01298661, dtype=float64)]
 
-PyTensor could provide shared JAX-like RNGs and allow RandomVariables to accept them,
-but that would break the spirit of one graph `->` multiple backends.
+``set_value`` is still observed, so reseeding with the same seed deterministically reproduces the original draws:
+
+>>> rng.set_value(np.random.default_rng(123))
+>>> print(jax_fn(), jax_fn())
+[Array(0.91459085, dtype=float64)] [Array(0.01298661, dtype=float64)]
+
+But ``get_value`` no longer observes the function's *own* updates. The JAX view advances privately and is never read back, so the host stays exactly at whatever was last set. Even after a call that advanced the JAX view, ``get_value`` still matches a freshly seeded seed-123 generator:
+
+>>> rng.set_value(np.random.default_rng(123))
+>>> _ = jax_fn()   # advances the private JAX view
+>>> rng.get_value().uniform() == np.random.default_rng(123).uniform()
+True
 
-Alternatively, PyTensor could try to use a more general type for RNGs that can be used across different backends,
-either directly or after some conversion operation (if such operations can be implemented in the different backends).
+With the default ``readback=True`` the same check would be ``False``, because the function's update is read back into the shared variable.

pytensor/compile/debug/debugmode.py

@@ -2000,6 +2000,7 @@ def __init__(
         name=None,
         no_fgraph_prep=False,
         trust_input=False,
+        readback=True,
     ):
         self.mode = mode
         self.profile = profile
@@ -2151,6 +2152,7 @@ def __init__(
         self.on_unused_input = on_unused_input  # Used for the pickling/copy
         self.name = name
         self.trust_input = trust_input
+        self.readback = readback
 
         self.required = [(i.value is None) for i in self.inputs]
         self.refeed = [

pytensor/compile/executor.py

@@ -17,6 +17,7 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import HasInnerGraph
 from pytensor.graph.utils import get_variable_trace_string
+from pytensor.link.backend_conversion import HOST, backend_handles
 from pytensor.link.basic import Container
 from pytensor.link.utils import raise_with_op
 
@@ -77,6 +78,8 @@ class Function:
         "_named_inputs",
         "_nodes_with_inner_function",
         "_potential_aliased_input_groups",
+        "_reconcile_shared",
+        "_shared_update_targets",
         "_update_input_storage",
         "input_storage",
         "maker",
@@ -242,6 +245,36 @@ def __init__(
                 )
             )
 
+        # Cross-backend shared coherence: implicit shareds of a Type some backend
+        # diverges on (e.g. RNGs under JAX) are reconciled before each call and
+        # re-stamped after an update, so a write by anyone reaches the next reader.
+        # Each is paired with this function's *effective* tag: its own backend if
+        # that backend diverges on the Type, else the host container it reads.
+        tag = maker.linker.backend_tag
+        divergent = [
+            inp
+            for inp in self.maker.expanded_inputs
+            if inp.implicit and inp.variable.type.is_backend_divergent
+        ]
+        if not maker.readback and tag != HOST:
+            # This function's backend updates are private; bind its views eager
+            # (host writes push into them) and enroll nothing, for a baseline __call__.
+            for inp in divergent:
+                if backend_handles(tag, inp.variable.type):
+                    inp.variable._view_storage(tag, eager=True)
+            self._reconcile_shared = ()
+            self._shared_update_targets = ()
+        else:
+            self._reconcile_shared = tuple(
+                (inp.variable, tag if backend_handles(tag, inp.variable.type) else HOST)
+                for inp in divergent
+            )
+            self._shared_update_targets = tuple(
+                (inp.variable, tag if backend_handles(tag, inp.variable.type) else HOST)
+                for inp in divergent
+                if inp.update is not None
+            )
+
         self._input_storage_data = tuple(
             container.storage for container in input_storage
         )
@@ -514,6 +547,7 @@ def checkSV(sv_ori, sv_rpl):
             accept_inplace=True,
             no_fgraph_prep=True,
             name=name,
+            readback=maker.readback,
         ).create(input_storage, storage_map=new_storage_map)
 
         for in_ori, in_cpy, ori, cpy in zip(
@@ -681,6 +715,10 @@ def __call__(self, *args, **kwargs):
         else:
             self._validate_inputs(args, kwargs)
 
+        if self._reconcile_shared:
+            for shared, tag in self._reconcile_shared:
+                shared._reconcile_into(tag)
+
         # Do the actual work
         try:
             if self.profile:
@@ -717,6 +755,9 @@ def __call__(self, *args, **kwargs):
         if self._has_updates:
             for i, input_storage in self._update_input_storage:
                 input_storage.storage[0] = outputs[i]
+            if self._shared_update_targets:
+                for shared, tag in self._shared_update_targets:
+                    shared._mark_written(tag)
             outputs = outputs[: self._n_returned_outputs]
 
         # Remove input and output values from storage data

pytensor/compile/maker.py

@@ -80,6 +80,7 @@ def function(
     profile: bool | ProfileStats | None = None,
     on_unused_input: str | None = None,
     trust_input: bool = False,
+    readback: bool = True,
 ):
     """
     Return a :class:`callable object <pytensor.compile.executor.Function>`
@@ -143,6 +144,13 @@ def function(
         that multiple inputs are not aliased to each other. Failure to meet any
         of these conditions can lead to computational errors or to the
         interpreter crashing.
+    readback: bool, default True
+        Only relevant for shared variables with a backend-specific representation
+        (e.g. random generators under the JAX backend). When True, updates this
+        function makes to such a shared are read back so other backends and
+        ``get_value`` observe them. Set False if this function's backend updates
+        are private (it only needs to observe host ``set_value``); this removes
+        all per-call cross-backend bookkeeping from the compiled function.
 
     Returns
     -------
@@ -257,6 +265,7 @@ def function(
         on_unused_input=on_unused_input,
         name=name,
         trust_input=trust_input,
+        readback=readback,
     )
     return m.create()
 
@@ -530,6 +539,7 @@ def __init__(
         name=None,
         no_fgraph_prep=False,
         trust_input=False,
+        readback=True,
     ):
         if profile:
             self._compile_start = time.perf_counter()
@@ -639,6 +649,7 @@ def __init__(
         self.on_unused_input = on_unused_input  # Used for the pickling/copy
         self.name = name
         self.trust_input = trust_input
+        self.readback = readback
         self.required = [(i.value is None) for i in self.inputs]
 
     def create(self, input_storage=None, storage_map=None):