diff --git a/doc/tutorial/prng.rst b/doc/tutorial/prng.rst index ed6b4b248e..bea8ac3d9e 100644 --- a/doc/tutorial/prng.rst +++ b/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,8 +532,8 @@ NumPy random generators can be natively used with the Numba backend. ├─ [] [id B] ├─ randomstate_rng [id C] ├─ NoneConst{None} [id D] - ├─ 0.0 [id E] - └─ 1.0 [id F] + ├─ 0.0 [id E] + └─ 1.0 [id F] [normal_rv{"(),()->()"}].0 [id A] 0 └─ ··· @@ -538,10 +541,10 @@ Inner graphs: [normal_rv{"(),()->()"}] [id A] ← normal_rv{"(),()->()"}.0 [id G] - ├─ *1- [id H] - ├─ *2- [id I] - ├─ *3- [id J] - └─ *4- [id K] + ├─ i1 [id H] + ├─ i2 [id I] + ├─ i3 [id J] + └─ i4 [id K] ← normal_rv{"(),()->()"}.1 [id G] └─ ··· @@ -553,22 +556,22 @@ Numba converts regular RandomVariable to RandomVariableWithCoreShape Ops, subtly >>> type(numba_fn.maker.fgraph.outputs[0].owner.op) -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] 0 - ├─ RNG() [id B] + ├─ rng [id B] ├─ NoneConst{None} [id C] ├─ 0.0 [id D] └─ 1.0 [id E] @@ -576,28 +579,40 @@ uniform_rv{"(),()->()"}.0 [id A] 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. diff --git a/pytensor/compile/debug/debugmode.py b/pytensor/compile/debug/debugmode.py index 135d8ecd83..993560e1c2 100644 --- a/pytensor/compile/debug/debugmode.py +++ b/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 = [ diff --git a/pytensor/compile/executor.py b/pytensor/compile/executor.py index 9de44ad9e3..64dbe6b264 100644 --- a/pytensor/compile/executor.py +++ b/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 diff --git a/pytensor/compile/maker.py b/pytensor/compile/maker.py index 7e9b103717..3c2931f386 100644 --- a/pytensor/compile/maker.py +++ b/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 ` @@ -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): diff --git a/pytensor/compile/sharedvalue.py b/pytensor/compile/sharedvalue.py index d466fd6036..c7b65e9605 100644 --- a/pytensor/compile/sharedvalue.py +++ b/pytensor/compile/sharedvalue.py @@ -8,6 +8,7 @@ from pytensor.graph.basic import Variable from pytensor.graph.utils import add_tag_trace +from pytensor.link.backend_conversion import HOST, backend_conversion from pytensor.link.basic import Container from pytensor.link.c.type import generic from pytensor.utils import add_lazy_dispatcher @@ -33,6 +34,42 @@ def collect_new_shareds(): __SHARED_CONTEXT__ = old_context +class _BackendState: + """Tracks backend-native views of a shared value, stored on its Container. + + Lives on the ``Container`` (shared across clones and across functions), not + the ``SharedVariable``. ``fresh`` is the set of tags whose copy currently + holds the authoritative value (``HOST`` = the canonical container); a write + by tag ``W`` leaves ``fresh == {W}``, and a read by tag ``R`` translates only + when ``R not in fresh``. So a single-backend loop never converts and never + writes state. ``eager`` tags are re-materialized on host writes rather than + invalidated (used by ``readback=False`` functions). + """ + + __slots__ = ("eager", "fresh", "views", "warned") + + def __init__(self): + self.views: dict[str, list] = {} + self.fresh: set[str] = {HOST} + self.eager: set[str] = set() + self.warned: set[str] = set() + + def reconcile_host(self, storage: list) -> str | None: + """Reconcile ``storage[0]`` from the authoritative backend view if stale. + + Returns the source tag when a (possibly lossy) conversion ran, else + ``None``. Silent: callers that want to warn inspect the return value. + """ + if HOST in self.fresh: + return None + source = next(iter(self.fresh)) + storage[0] = backend_conversion(source).from_native( # type: ignore[union-attr] + self.views[source][0] + ) + self.fresh.add(HOST) + return source + + class SharedVariable(Variable): """Variable that is shared between compiled functions.""" @@ -94,6 +131,90 @@ def __init__( self._default_update: Variable | None = None + def _backend_state(self) -> _BackendState: + state = self.container._backend_state + if state is None: + state = self.container._backend_state = _BackendState() + return state + + def _host_value(self): + """Return the current host value, reconciling from a fresh backend if stale. + + Warns once per source when the fresh backend's representation is lossy, + so the reconciled host value does not continue that backend's stream. + """ + state = self.container._backend_state + if state is None: + return self.container.value + source = state.reconcile_host(self.container.storage) + if ( + source is not None + and backend_conversion(source).lossy + and source not in state.warned + ): + warnings.warn( + f"{self} was last advanced by the {source!r} backend; converting " + f"its state back is lossy, so the reconciled value does not " + f"continue the {source!r} stream.", + UserWarning, + stacklevel=3, + ) + state.warned.add(source) + return self.container.value + + def _view_storage(self, backend: str, *, eager: bool = False) -> list: + """Return (materializing if needed) the native storage list for ``backend``. + + Pass ``eager=True`` to have host writes re-materialize this view eagerly + (see ``readback=False``) instead of invalidating it lazily. + """ + state = self._backend_state() + if eager: + state.eager.add(backend) + storage = state.views.get(backend) + if storage is None: + native = backend_conversion(backend).to_native( # type: ignore[union-attr] + self._host_value() + ) + storage = state.views[backend] = [native] + state.fresh.add(backend) + return storage + + def _mark_written(self, backend: str) -> None: + """Record that ``backend`` produced the current authoritative value. + + Invalidates every other copy. A no-op (no allocation, no write) when + ``backend`` is already the sole fresh copy, so a steady loop pays nothing. + """ + state = self.container._backend_state + if state is None: + return + fresh = state.fresh + if len(fresh) != 1 or backend not in fresh: + state.fresh = {backend} + + def _reconcile_into(self, backend: str) -> None: + """Bring ``backend``'s copy up to the authoritative value, if it is stale. + + A no-op (single membership check) unless another backend wrote since this + one last synced, so a single-backend loop never converts. + """ + state = self.container._backend_state + if state is None or backend in state.fresh: + return + if backend == HOST: + self._host_value() + return + native = backend_conversion(backend).to_native( # type: ignore[union-attr] + self._host_value() + ) + storage = state.views.get(backend) + if storage is None: + state.views[backend] = [native] + else: + storage[0] = native + state.fresh.add(backend) + def get_value(self, borrow=False, return_internal_type=False): """ Get the non-symbolic value associated with this SharedVariable. @@ -106,16 +227,17 @@ def get_value(self, borrow=False, return_internal_type=False): True to permit the returning of an arbitrary type object used internally to store the shared variable. + The host value is reconciled first if a compiled function on a backend + with its own representation (e.g. JAX) advanced it more recently. + Only with borrow=False and return_internal_type=True does this function guarantee that you actually get the internal object. But in that case, you may get different return types when using different compute devices. """ - if borrow: - return self.container.value - else: - return copy.deepcopy(self.container.value) + value = self._host_value() + return value if borrow else copy.deepcopy(value) def set_value(self, new_value, borrow=False): """ @@ -134,6 +256,14 @@ def set_value(self, new_value, borrow=False): self.container.value = new_value else: self.container.value = copy.deepcopy(new_value) + # Invalidate the cached backend views. Eager ones (readback=False, which + # skip the reconcile bracket) are refreshed now; the rest lazily on use. + state = self.container._backend_state + if state is not None: + host = self.container.value + for tag in state.eager: + state.views[tag][0] = backend_conversion(tag).to_native(host) + state.fresh = {HOST, *state.eager} def clone(self, **kwargs): name = kwargs.get("name", self.name) diff --git a/pytensor/graph/type.py b/pytensor/graph/type.py index f165d2f258..e07f0a6138 100644 --- a/pytensor/graph/type.py +++ b/pytensor/graph/type.py @@ -33,6 +33,15 @@ class Type(MetaObject, Generic[D]): # noqa: UP046 The `Type` that will be created by a call to `Type.make_constant`. """ + is_backend_divergent: bool = False + """ + Whether this `Type`'s concrete value has a backend-specific representation + (e.g. a random generator or sparse matrix that some backends store + differently from the host). Shared variables of such a type are tracked for + cross-backend value reconciliation. This is intrinsic to the type and does + not depend on which backend dispatch modules have been imported. + """ + def in_same_class(self, otype: "Type") -> bool | None: """Determine if another `Type` represents a subset from the same "class" of types represented by `self`. diff --git a/pytensor/link/backend_conversion.py b/pytensor/link/backend_conversion.py new file mode 100644 index 0000000000..6d8a331e1e --- /dev/null +++ b/pytensor/link/backend_conversion.py @@ -0,0 +1,49 @@ +"""Cross-backend conversion registry for shared-variable representations. + +Some backends (e.g. JAX) cannot operate on the host representation of certain +shared variables (a numpy ``Generator``, a scipy sparse matrix) and require a +backend-native representation instead. A backend registers how to convert a host +value *to* its native form and, when possible, back *from* it. + +A compiled function keeps a per-backend view of such a shared value and +reconciles it lazily, so a single-backend loop never pays a conversion: the +conversion runs only when a value crosses backends (see +`SharedVariable._reconcile_into`). +""" + +from collections.abc import Callable +from dataclasses import dataclass + + +HOST = "host" + + +@dataclass(frozen=True) +class BackendConversion: + tag: str + handles: Callable[[object], bool] + to_native: Callable[[object], object] + from_native: Callable[[object], object] + lossy: bool = False + + +_CONVERSIONS: dict[str, BackendConversion] = {} + + +def register_backend_conversion(conversion: BackendConversion) -> None: + _CONVERSIONS[conversion.tag] = conversion + + +def backend_conversion(tag: str) -> BackendConversion | None: + return _CONVERSIONS.get(tag) + + +def backend_handles(tag: str, type_) -> bool: + """Whether ``tag`` needs a native view for this Type (a type-and-backend property). + + Whether the Type is divergent for *some* backend at all is the intrinsic + ``Type.is_backend_divergent`` property (import-order-independent); this asks + whether the specific, already-imported ``tag`` provides a conversion for it. + """ + conversion = _CONVERSIONS.get(tag) + return conversion is not None and conversion.handles(type_) diff --git a/pytensor/link/basic.py b/pytensor/link/basic.py index 80e57afff1..f99fb042ac 100644 --- a/pytensor/link/basic.py +++ b/pytensor/link/basic.py @@ -7,6 +7,7 @@ from pytensor.graph.basic import Apply, Variable from pytensor.graph.fg import FunctionGraph from pytensor.graph.type import Type +from pytensor.link.backend_conversion import HOST, backend_handles from pytensor.link.utils import gc_helper, map_storage, raise_with_op, streamline from pytensor.utils import difference @@ -15,6 +16,7 @@ from numpy.typing import NDArray from pytensor.compile.debug.profiling import ProfileStats + from pytensor.compile.sharedvalue import _BackendState from pytensor.graph.op import ( BasicThunkType, InputStorageType, @@ -79,6 +81,9 @@ def __init__( self.readonly = readonly self.strict = strict self.allow_downcast = allow_downcast + # Lazily holds backend-native views of a shared value (see backend_conversion). + # Stays None for the vast majority of Containers, which need no conversion. + self._backend_state: _BackendState | None = None def __get__(self) -> Any: return self.storage[0] @@ -113,6 +118,10 @@ def __repr__(self): return "<" + repr(self.storage[0]) + ">" def __deepcopy__(self, memo: dict[int, Any]) -> "Container": + # If a backend advanced the value more recently than the host, bring the + # host copy up to date first so the deep copy doesn't capture a stale value. + if self._backend_state is not None: + self._backend_state.reconcile_host(self.storage) data_was_in_memo = id(self.storage[0]) in memo r = type(self)( deepcopy(self.type, memo=memo), @@ -153,6 +162,10 @@ class Linker(ABC): required_rewrites: tuple[str, ...] = ("minimum_compile",) incompatible_rewrites: tuple[str, ...] = () + # Identifies the backend representation this linker consumes. Backends whose + # native format matches the host (C, numba, python) keep HOST; those that + # diverge (JAX) override it and register a BackendConversion under this tag. + backend_tag: str = HOST def __init__( self, @@ -181,6 +194,38 @@ def clone(self, allow_gc: bool | None = None) -> "Linker": new._allow_gc = allow_gc return new + def _bind_shared_backend_storage(self, input_storage, storage_map): + """Point divergent shared inputs at their backend-native view storage. + + A shared variable whose Type this backend can't consume in host format + has the thunks' storage (in both ``storage_map`` and ``input_storage``) + repointed at a per-backend view materialized on its Container. The graph + keeps the original variable, so ``set_value``/``swap`` continue to target + storage the thunks read. A no-op for host backends, so any linker can call + it after ``map_storage``. + """ + # Local import to avoid a circular import: sharedvalue imports this module. + from pytensor.compile.sharedvalue import SharedVariable + + tag = self.backend_tag + if tag == HOST: + return + for inp in self.fgraph.inputs: + if not isinstance(inp, SharedVariable) or not backend_handles( + tag, inp.type + ): + continue + old = storage_map[inp] + view = inp._view_storage(tag) + if old is view: + continue + storage_map[inp] = view + for idx, cell in enumerate(input_storage): + # Identity, because input_storage may contain numpy arrays (issue #314) + if cell is old: + input_storage[idx] = view + break + @abstractmethod def make_thunk( self, **kwargs @@ -329,6 +374,7 @@ def make_all( input_storage, output_storage, storage_map = map_storage( fgraph, order, input_storage, output_storage, storage_map ) + self._bind_shared_backend_storage(input_storage, storage_map) compute_map = {} for k in storage_map: @@ -712,6 +758,7 @@ def make_all(self, input_storage=None, output_storage=None, storage_map=None): input_storage, output_storage, storage_map = map_storage( fgraph, nodes, input_storage, output_storage, storage_map ) + self._bind_shared_backend_storage(input_storage, storage_map) compute_map = {} for k in storage_map: diff --git a/pytensor/link/jax/dispatch/random.py b/pytensor/link/jax/dispatch/random.py index edd9fff8b5..18c75c652f 100644 --- a/pytensor/link/jax/dispatch/random.py +++ b/pytensor/link/jax/dispatch/random.py @@ -4,14 +4,16 @@ import jax.numpy as jnp import numpy as np from numpy.random import Generator -from numpy.random.bit_generator import ( # type: ignore[attr-defined] - _coerce_to_uint32_array, -) import pytensor.tensor.random.basic as ptr from pytensor.graph import Constant +from pytensor.link.backend_conversion import ( + BackendConversion, + register_backend_conversion, +) 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 from pytensor.tensor.type_other import NoneTypeT @@ -23,9 +25,6 @@ except ImportError: numpyro_available = False -numpy_bit_gens = {"MT19937": 0, "PCG64": 1, "Philox": 2, "SFC64": 3} - - SIZE_NOT_COMPATIBLE = """JAX random variables require concrete values for the `size` parameter of the distributions. Concrete values are either constants: @@ -57,27 +56,43 @@ def assert_size_argument_jax_compatible(node): @jax_typify.register(Generator) def jax_typify_Generator(rng, **kwargs): - state = rng.bit_generator.state - state["bit_generator"] = numpy_bit_gens[state["bit_generator"]] - - # XXX: Is this a reasonable approach? - state["jax_state"] = _coerce_to_uint32_array(state["state"]["state"])[0:2] - - # The "state" and "inc" values in a NumPy `Generator` are 128 bits, which - # JAX can't handle, so we split these values into arrays of 32 bit integers - # and then combine the first two into a single 64 bit integers. - # - # XXX: Depending on how we expect these values to be used, is this approach - # reasonable? - # - # TODO: We might as well remove these altogether, since this conversion - # should only occur once (e.g. when the graph is converted/JAX-compiled), - # and, from then on, we use the custom "jax_state" value. - inc_32 = _coerce_to_uint32_array(state["state"]["inc"]) - 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 + """Derive a JAX threefry key from a host numpy ``Generator``. + + The forward half of the lossy map completed by `jax_detypify_Generator`. + Hashing the ~256-bit PCG64 state into 2 words is not reversible, so this + seeds the JAX key deterministically but does not preserve the host stream. + """ + # A JAX threefry key holds only 2 uint32 words, far fewer than the ~256 bits + # of a NumPy Generator's state. Rather than truncate to the low + # (worst-quality) bits of the LCG state and drop the stream `inc`, fold the + # whole state through SeedSequence's mixer so all the entropy is hashed into + # the 2 words we keep. + state = rng.bit_generator.state["state"] + seed = np.random.SeedSequence([state["state"], state["inc"]]) + return seed.generate_state(2, dtype=np.uint32) + + +def jax_detypify_Generator(key): + """Reconcile a JAX-advanced threefry key back to a host numpy ``Generator``. + + The reverse half of the lossy map started by `jax_typify_Generator`. The + host cannot continue the JAX stream, so this reseeds a ``Generator`` from the + key bits deterministically, starting a new, unrelated stream. + """ + key = np.asarray(key).astype(np.uint32).ravel() + seed = np.random.SeedSequence(int.from_bytes(key.tobytes(), "little")) + return np.random.Generator(np.random.PCG64(seed)) + + +register_backend_conversion( + BackendConversion( + tag="jax", + handles=lambda type_: isinstance(type_, RandomType), + to_native=jax_typify, + from_native=jax_detypify_Generator, + lossy=True, + ) +) @jax_funcify.register(ptr.RandomVariable) @@ -106,24 +121,20 @@ def jax_funcify_RandomVariable(op: ptr.RandomVariable, node, **kwargs): assert_size_argument_jax_compatible(node) def sample_fn(rng, size, *parameters): - rng_key = rng["jax_state"] - rng_key, sampling_key = jax.random.split(rng_key, 2) - rng["jax_state"] = rng_key + next_rng, sampling_key = jax.random.split(rng, 2) sample = jax_sample_fn(op, node=node)( sampling_key, size, out_dtype, *parameters ) - return (rng, sample) + return (next_rng, sample) else: def sample_fn(rng, size, *parameters): - rng_key = rng["jax_state"] - rng_key, sampling_key = jax.random.split(rng_key, 2) - rng["jax_state"] = rng_key + next_rng, sampling_key = jax.random.split(rng, 2) sample = jax_sample_fn(op, node=node)( sampling_key, static_size, out_dtype, *parameters ) - return (rng, sample) + return (next_rng, sample) return sample_fn diff --git a/pytensor/link/jax/linker.py b/pytensor/link/jax/linker.py index 4ab0224881..6d05bc6a74 100644 --- a/pytensor/link/jax/linker.py +++ b/pytensor/link/jax/linker.py @@ -1,14 +1,11 @@ -import warnings - -from numpy.random import Generator - -from pytensor.compile.sharedvalue import SharedVariable, shared from pytensor.link.basic import JITLinker class JAXLinker(JITLinker): """A `Linker` that JIT-compiles NumPy-based operations using JAX.""" + backend_tag = "jax" + required_rewrites = ( "minimum_compile", "jax", @@ -34,58 +31,6 @@ def __init__(self, *args, **kwargs): def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs): from pytensor.link.jax.dispatch import jax_funcify from pytensor.link.jax.dispatch.shape import JAXShapeTuple - from pytensor.tensor.random.type import RandomType - - shared_rng_inputs = [ - inp - for inp in fgraph.inputs - if (isinstance(inp, SharedVariable) and isinstance(inp.type, RandomType)) - ] - - # Replace any shared RNG inputs so that their values can be updated in place - # without affecting the original RNG container. This is necessary because - # JAX does not accept Generators as inputs, and they will have to - # be tipyfied - if shared_rng_inputs: - warnings.warn( - f"The RandomType SharedVariables {shared_rng_inputs} will not be used " - f"in the compiled JAX graph. Instead a copy will be used.", - UserWarning, - ) - new_shared_rng_inputs = [ - shared(inp.get_value(borrow=False)) for inp in shared_rng_inputs - ] - - fgraph.replace_all( - zip(shared_rng_inputs, new_shared_rng_inputs, strict=True), - import_missing=True, - reason="JAXLinker.fgraph_convert", - ) - - for old_inp, new_inp in zip( - shared_rng_inputs, new_shared_rng_inputs, strict=True - ): - new_inp_storage = [new_inp.get_value(borrow=True)] - storage_map[new_inp] = new_inp_storage - old_inp_storage = storage_map.pop(old_inp) - # Find index of old_inp_storage in input_storage - for input_storage_idx, input_storage_item in enumerate(input_storage): - # We have to establish equality based on identity because input_storage may contain numpy arrays - if input_storage_item is old_inp_storage: - break - else: # no break - raise ValueError() - input_storage[input_storage_idx] = new_inp_storage - # We need to change the order of the inputs of the FunctionGraph - # so that the new input is in the same position as to old one, - # to align with the storage_map. We hope this is safe! - old_inp_fgrap_index = fgraph.inputs.index(old_inp) - fgraph.remove_input( - old_inp_fgrap_index, - reason="JAXLinker.fgraph_convert", - ) - fgraph.inputs.remove(new_inp) - fgraph.inputs.insert(old_inp_fgrap_index, new_inp) fgraph_inputs = fgraph.inputs clients = fgraph.clients @@ -129,14 +74,4 @@ def convert_scalar_shape_inputs( return convert_scalar_shape_inputs def create_thunk_inputs(self, storage_map): - from pytensor.link.jax.dispatch import jax_typify - - thunk_inputs = [] - for n in self.fgraph.inputs: - sinput = storage_map[n] - if isinstance(sinput[0], Generator): - # Neet to convert Generator into JAX PRNGkey - sinput[0] = jax_typify(sinput[0]) - thunk_inputs.append(sinput) - - return thunk_inputs + return [storage_map[n] for n in self.fgraph.inputs] diff --git a/pytensor/link/vm.py b/pytensor/link/vm.py index 239f73df80..de73dd6eab 100644 --- a/pytensor/link/vm.py +++ b/pytensor/link/vm.py @@ -1216,6 +1216,7 @@ def make_all( input_storage, output_storage, storage_map = map_storage( fgraph, order, input_storage, output_storage, storage_map ) + self._bind_shared_backend_storage(input_storage, storage_map) compute_map = {} for k in storage_map: compute_map[k] = [k.owner is None] diff --git a/pytensor/tensor/random/type.py b/pytensor/tensor/random/type.py index 5ff5354d31..658ebcd88e 100644 --- a/pytensor/tensor/random/type.py +++ b/pytensor/tensor/random/type.py @@ -10,23 +10,11 @@ T = TypeVar("T") -gen_states_keys = { - "MT19937": (["state"], ["key", "pos"]), - "PCG64": (["state", "has_uint32", "uinteger"], ["state", "inc"]), - "Philox": ( - ["state", "buffer", "buffer_pos", "has_uint32", "uinteger"], - ["counter", "key"], - ), - "SFC64": (["state", "has_uint32", "uinteger"], ["state"]), -} - -# We map bit generators to an integer index so that we can avoid using strings -numpy_bit_gens = {0: "MT19937", 1: "PCG64", 2: "Philox", 3: "SFC64"} - - class RandomType(Type[T]): r"""A Type wrapper for `numpy.random.Generator.""" + is_backend_divergent = True + class AbstractRandomGeneratorType(RandomType[Generator]): r"""Abstract base for random generator types wrapping `numpy.random.Generator`. @@ -54,34 +42,12 @@ def filter(self, data, strict=False, allow_downcast=None): if isinstance(data, Generator): return data - if not strict and isinstance(data, dict): - if "bit_generator" not in data: - raise TypeError() - else: - bit_gen_key = data["bit_generator"] - - if hasattr(bit_gen_key, "_value"): - bit_gen_key = int(bit_gen_key._value) - bit_gen_key = numpy_bit_gens[bit_gen_key] - - gen_keys, state_keys = gen_states_keys[bit_gen_key] - - for key in gen_keys: - if key not in data: - raise TypeError() - - for key in state_keys: - if key not in data["state"]: - raise TypeError() - - return data - raise TypeError() @staticmethod def values_eq(a, b): - sa = a if isinstance(a, dict) else a.bit_generator.state - sb = b if isinstance(b, dict) else b.bit_generator.state + sa = a.bit_generator.state + sb = b.bit_generator.state def _eq(sa, sb): for key in sa: @@ -107,10 +73,6 @@ class RandomGeneratorType(AbstractRandomGeneratorType): `Generator` objects that would appear to be equal do not compare equal with the ``==`` operator. - This `Type` also works with a ``dict`` derived from - `Generator.__get_state__`, unless the ``strict`` argument to `Type.filter` - is explicitly set to ``True``. - """ def __repr__(self): diff --git a/tests/link/jax/test_random.py b/tests/link/jax/test_random.py index 53f91d9d34..794709cafc 100644 --- a/tests/link/jax/test_random.py +++ b/tests/link/jax/test_random.py @@ -26,13 +26,6 @@ from pytensor.link.jax.dispatch.random import numpyro_available -def compile_random_function(*args, mode=jax_mode, **kwargs): - with pytest.warns( - UserWarning, match=r"The RandomType SharedVariables \[.+\] will not be used" - ): - return function(*args, mode=mode, **kwargs) - - def test_random_RandomStream(): """Two successive calls of a compiled graph using `RandomStream` should return different values. @@ -41,7 +34,7 @@ def test_random_RandomStream(): srng = RandomStream(seed=123) out = srng.normal() - srng.normal() - fn = compile_random_function([], out) + fn = function([], out, mode=jax_mode) jax_res_1 = fn() jax_res_2 = fn() @@ -54,10 +47,9 @@ 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}) - assert f() != f() + f = function([], [x], updates={rng: next_rng}, mode=jax_mode) - # Check that original rng variable content was not overwritten when calling jax_typify + # Materializing the JAX view must not corrupt the host Generator assert all( a == b if not isinstance(a, np.ndarray) else np.array_equal(a, b) for a, b in zip( @@ -67,6 +59,50 @@ def test_random_updates(rng_ctor): ) ) + assert f() != f() + + +def test_rng_shared_across_backends(): + """A JAX and a host function sharing one RNG advance a single stream: a call + on either backend is seen by the next call on the other.""" + rng = shared(np.random.default_rng(0), borrow=False) + next_rng, x = rng.normal() + jax_fn = function([], x, updates={rng: next_rng}, mode=jax_mode) + host_fn = function([], x, updates={rng: next_rng}) # C backend + + # A JAX call advances the shared stream, so the next host draw differs from + # the host draw taken straight from the same seed. + rng.set_value(np.random.default_rng(0)) + host_only = host_fn() + rng.set_value(np.random.default_rng(0)) + jax_fn() + assert host_fn() != host_only + + # And symmetrically, a host call is seen by the next JAX draw. + rng.set_value(np.random.default_rng(0)) + jax_only = jax_fn() + rng.set_value(np.random.default_rng(0)) + host_fn() + assert jax_fn() != jax_only + + +def test_rng_readback_false(): + """readback=False advances the RNG privately: host get_value is unaffected by + the function's updates, but the function still observes user set_value.""" + rng = shared(np.random.default_rng(0), borrow=False) + next_rng, x = rng.normal() + fn = function([], x, updates={rng: next_rng}, mode=jax_mode, readback=False) + + first = fn() + assert fn() != first # advances across calls + + host_state = rng.get_value().bit_generator.state + fn() + assert rng.get_value().bit_generator.state == host_state # host untouched + + rng.set_value(np.random.default_rng(0)) + assert fn() == first # user set_value is observed + @pytest.mark.parametrize("noise_first", (False, True)) def test_replaced_shared_rng_storage_order(noise_first): @@ -86,7 +122,7 @@ def test_replaced_shared_rng_storage_order(noise_first): mu: pt.grad(out, mu), rng: next_rng, } - f = compile_random_function([], [out], updates=updates, mode="JAX") + f = function([], [out], updates=updates, mode="JAX") # The bug was found when noise used to be the first input of the fgraph # If this changes, the test may need to be tweaked to keep the save coverage @@ -126,7 +162,7 @@ def test_replaced_shared_rng_storage_ordering_equality(): # This function replaces inp by input_shared in the update expression # This is what caused the RNG to appear later than inp_shared in the input_storage - fn = compile_random_function( + fn = function( inputs=[], outputs=[], updates={inp_shared: inp_update}, @@ -482,7 +518,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) + g_fn = function(dist_params, g, mode=jax_mode) samples = g_fn(*test_values) bcast_dist_args = np.broadcast_arrays(*test_values) @@ -528,7 +564,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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() np.testing.assert_allclose(samples.mean(axis=0), 0.5, 1) @@ -539,7 +575,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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() np.testing.assert_allclose(samples.mean(axis=0), mu, atol=0.1) @@ -559,7 +595,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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() np.testing.assert_allclose(samples.mean(axis=0), 0.5, 1) @@ -568,7 +604,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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() assert samples.shape == (10_000,) # Elements are picked at equal frequency @@ -577,7 +613,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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() assert samples.shape == (5, 2) # Only odd numbers are picked @@ -586,7 +622,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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() assert samples.shape == (2, 49) # Elements are unique @@ -601,7 +637,7 @@ def test_random_choice(): rng=rng, replace=False, ) - g_fn = compile_random_function([], g) + g_fn = function([], g, mode=jax_mode) samples = g_fn() assert samples.shape == (3,) # Elements are unique @@ -613,14 +649,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) + g_fn = function([], g, mode=jax_mode) 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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() assert samples.shape == (1000,) assert np.all(samples % 2 == 1) @@ -630,7 +666,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) + g_fn = function([], g, mode=jax_mode) permuted = g_fn() with pytest.raises(AssertionError): np.testing.assert_allclose(array, permuted) @@ -653,7 +689,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) + g_fn = function([], g, mode=jax_mode) 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) @@ -664,7 +700,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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() np.testing.assert_allclose(samples.mean(axis=0), n * (1 - p) / p, rtol=0.1) np.testing.assert_allclose( @@ -678,7 +714,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) + g_fn = function([], g, mode=jax_mode) 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) @@ -693,7 +729,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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() np.testing.assert_allclose(samples.mean(axis=0), n * a / (a + b), rtol=0.1) np.testing.assert_allclose( @@ -712,7 +748,7 @@ def test_multinomial(): size = (10_000, 2) g = pt.random.multinomial(n, p, size=size, rng=rng) - g_fn = compile_random_function([], g, mode="JAX") + g_fn = function([], g, mode="JAX") samples = g_fn() np.testing.assert_allclose(samples.mean(axis=0), n[..., None] * p, rtol=0.1) np.testing.assert_allclose( @@ -726,7 +762,7 @@ def test_multinomial(): pt_p = pt.matrix("p") g = pt.random.multinomial(pt_n, pt_p, rng=rng, size=None) - g_fn = compile_random_function([pt_n, pt_p], g, mode="JAX") + g_fn = function([pt_n, pt_p], g, mode="JAX") samples = g_fn(n, p) np.testing.assert_allclose(samples.mean(axis=0), n[0, :, None] * p, rtol=0.1) np.testing.assert_allclose( @@ -735,13 +771,13 @@ def test_multinomial(): # Test with p=0 g = pt.random.multinomial(n=5, p=pt.eye(4)) - g_fn = compile_random_function([], g, mode="JAX") + g_fn = function([], g, mode="JAX") samples = g_fn() np.testing.assert_array_equal(samples, np.eye(4) * 5) # Test with n=0 g = pt.random.multinomial(n=0, p=np.ones(4) / 4) - g_fn = compile_random_function([], g, mode="JAX") + g_fn = function([], g, mode="JAX") samples = g_fn() np.testing.assert_array_equal(samples, np.zeros(4)) @@ -754,7 +790,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) + g_fn = function([], g, mode=jax_mode) samples = g_fn() np.testing.assert_allclose( samples.mean(axis=0), (mu + np.pi) % (2.0 * np.pi) - np.pi, rtol=0.1 @@ -798,10 +834,7 @@ def rng_fn(cls, rng, size): out = nonexistentrv(rng=rng) with pytest.raises(NotImplementedError): - with pytest.warns( - UserWarning, match=r"The RandomType SharedVariables \[.+\] will not be used" - ): - compare_jax_and_py([], [out], []) + compare_jax_and_py([], [out], []) def test_random_custom_implementation(): @@ -830,10 +863,7 @@ def sample_fn(rng, size, dtype, *parameters): nonexistentrv = CustomRV() rng = shared(np.random.default_rng(123)) out = nonexistentrv(rng=rng) - with pytest.warns( - UserWarning, match=r"The RandomType SharedVariables \[.+\] will not be used" - ): - compare_jax_and_py([], [out], []) + compare_jax_and_py([], [out], []) class TestRandomShapeInputs: @@ -851,14 +881,14 @@ def test_random_concrete_shape(self): 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) + jax_fn = function([x_pt], out, mode=jax_mode) assert jax_fn(np.ones((2, 3))).shape == (2, 3) def test_random_concrete_shape_from_param(self): 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) + jax_fn = function([x_pt], out, mode=jax_mode) assert jax_fn(np.ones((2, 3))).shape == (2, 3) def test_random_concrete_shape_subtensor(self): @@ -876,7 +906,7 @@ def test_random_concrete_shape_subtensor(self): 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) + jax_fn = function([x_pt], out, mode=jax_mode) assert jax_fn(np.ones((2, 3))).shape == (3,) def test_random_concrete_shape_subtensor_tuple(self): @@ -891,7 +921,7 @@ def test_random_concrete_shape_subtensor_tuple(self): 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) + jax_fn = function([x_pt], out, mode=jax_mode) assert jax_fn(np.ones((2, 3))).shape == (2,) def test_random_scalar_shape_input(self): @@ -899,12 +929,12 @@ def test_random_scalar_shape_input(self): dim1 = pt.scalar("dim1", dtype=int) out = pt.random.normal(0, 1, size=dim0) - jax_fn = compile_random_function([dim0], out) + jax_fn = function([dim0], out, mode=jax_mode) assert jax_fn(np.array(2)).shape == (2,) assert jax_fn(np.array(3)).shape == (3,) out = pt.random.normal(0, 1, size=[dim0, dim1]) - jax_fn = compile_random_function([dim0, dim1], out) + jax_fn = function([dim0, dim1], out, mode=jax_mode) assert jax_fn(np.array(2), np.array(3)).shape == (2, 3) assert jax_fn(np.array(4), np.array(5)).shape == (4, 5) @@ -916,7 +946,7 @@ def test_random_scalar_shape_input_not_supported(self): out1 = pt.random.normal(0, 1, size=dim) # An operation that wouldn't work if we replaced 0d array by integer out2 = dim[...].set(1) - jax_fn = compile_random_function([dim], [out1, out2]) + jax_fn = function([dim], [out1, out2], mode=jax_mode) res1, res2 = jax_fn(np.array(2)) assert res1.shape == (2,) @@ -930,7 +960,7 @@ def test_random_scalar_shape_input_not_supported2(self): # This could theoretically be supported # but would require knowing that * 2 is a safe operation for a python integer out = pt.random.normal(0, 1, size=dim * 2) - jax_fn = compile_random_function([dim], out) + jax_fn = function([dim], out, mode=jax_mode) assert jax_fn(np.array(2)).shape == (4,) @pytest.mark.xfail( @@ -940,7 +970,7 @@ def test_random_vector_shape_graph_input(self): shape = pt.vector("shape", shape=(2,), dtype=int) out = pt.random.normal(0, 1, size=shape) - jax_fn = compile_random_function([shape], out) + jax_fn = function([shape], out, mode=jax_mode) assert jax_fn(np.array([2, 3])).shape == (2, 3) assert jax_fn(np.array([4, 5])).shape == (4, 5) @@ -950,16 +980,16 @@ def test_constant_shape_after_graph_rewriting(self): assert x.type.shape == (None, None) with pytest.raises(TypeError): - compile_random_function([size], x)([2, 5]) + function([size], x, mode=jax_mode)([2, 5]) # Rebuild with strict=True so output type is not updated # This reflects cases where size is constant folded during rewrites but the RV node is not recreated new_x = clone_replace(x, {size: pt.constant([2, 5])}, rebuild_strict=True) assert new_x.type.shape == (None, None) - assert compile_random_function([], new_x)().shape == (2, 5) + assert function([], new_x, mode=jax_mode)().shape == (2, 5) # Rebuild with strict=False, so output type is updated # This uses a different path in the dispatch implementation new_x = clone_replace(x, {size: pt.constant([2, 5])}, rebuild_strict=False) assert new_x.type.shape == (2, 5) - assert compile_random_function([], new_x)().shape == (2, 5) + assert function([], new_x, mode=jax_mode)().shape == (2, 5) diff --git a/tests/link/jax/test_scan.py b/tests/link/jax/test_scan.py index b7784c7285..bc98cf5acd 100644 --- a/tests/link/jax/test_scan.py +++ b/tests/link/jax/test_scan.py @@ -1,5 +1,3 @@ -import re - import numpy as np import pytest @@ -167,11 +165,7 @@ def update_fn(rng): ) # Without updates - with pytest.warns( - UserWarning, - match=re.escape("[rng] will not be used in the compiled JAX graph"), - ): - jax_fn = function([], [xs], updates=None, mode="JAX") + jax_fn = function([], [xs], updates=None, mode="JAX") res1, res2 = jax_fn(), jax_fn() assert np.unique(res1).size == 10 @@ -179,11 +173,7 @@ def update_fn(rng): np.testing.assert_array_equal(res1, res2) # With updates - with pytest.warns( - UserWarning, - match=re.escape("[rng] will not be used in the compiled JAX graph"), - ): - jax_fn = function([], [xs], updates=update, mode="JAX") + jax_fn = function([], [xs], updates=update, mode="JAX") res1, res2 = jax_fn(), jax_fn() assert np.unique(res1).size == 10 diff --git a/tests/link/test_backend_conversion.py b/tests/link/test_backend_conversion.py new file mode 100644 index 0000000000..b411c98ab4 --- /dev/null +++ b/tests/link/test_backend_conversion.py @@ -0,0 +1,158 @@ +"""Backend-agnostic tests for cross-backend shared-value coherence. + +A minimal backend (`CounterLinker`) whose native representation of a divergent +`CounterType` differs from the host stands in for a real backend such as JAX. +Compiling and calling real functions on it exercises the actual path -- the +reconcile/mark bracket in `Function.__call__` and the storage binding in +`Linker._bind_shared_backend_storage` -- so the behaviour we care about (a write +by either side reaching the next read on the other) is tested end-to-end without +depending on JAX. `CounterLinker` is a plain (non-JIT) `PerformLinker`, which +also shows the divergence machinery is not tied to JIT backends. +""" + +import copy +import warnings + +import numpy as np +import pytest + +import pytensor +from pytensor.compile.mode import Mode +from pytensor.compile.sharedvalue import SharedVariable, shared +from pytensor.graph.basic import Apply +from pytensor.graph.op import Op +from pytensor.graph.rewriting.db import RewriteDatabaseQuery +from pytensor.graph.type import Type +from pytensor.link.backend_conversion import ( + _CONVERSIONS, + BackendConversion, + register_backend_conversion, +) +from pytensor.link.basic import PerformLinker +from pytensor.tensor.type import TensorType + + +class CounterType(Type): + """Divergent scalar type: the host holds a python int, the backend a dict.""" + + __props__ = () + is_backend_divergent = True + + def filter(self, data, strict=False, allow_downcast=None): + return int(data) + + +counter_type = CounterType() +_int_scalar = TensorType("int64", ()) + + +class Step(Op): + """(counter) -> (next_counter, value): read the counter, return it, advance it.""" + + __props__ = () + + def make_node(self, c): + return Apply(self, [c], [counter_type(), _int_scalar()]) + + def perform(self, node, inputs, outputs): + (c,) = inputs + # The op runs on the backend-native representation (bound in by the + # linker), not the host int -- proof the storage binding took effect. + assert isinstance(c, dict) + n = c["count"] + outputs[0][0] = {"count": n + 1} + outputs[1][0] = np.int64(n) + + +class CounterLinker(PerformLinker): + """An op-by-op backend that runs `perform` on the native representation.""" + + def __init__(self, tag, **kwargs): + self.backend_tag = tag + super().__init__(**kwargs) + + +def _make_backend(tag, *, lossy=False): + register_backend_conversion( + BackendConversion( + tag=tag, + handles=lambda t: isinstance(t, CounterType), + to_native=lambda h: {"count": int(h)}, + from_native=lambda n: n["count"], + lossy=lossy, + ) + ) + return Mode(linker=CounterLinker(tag), optimizer=RewriteDatabaseQuery(include=[])) + + +@pytest.fixture +def counter_backend(): + tag = "counter" + yield _make_backend(tag) + _CONVERSIONS.pop(tag, None) + + +@pytest.fixture +def lossy_backend(): + tag = "counter_lossy" + yield _make_backend(tag, lossy=True) + _CONVERSIONS.pop(tag, None) + + +def _counter_function(mode): + s = SharedVariable(type=counter_type, value=0, strict=False, name="c") + next_c, value = Step()(s) + return s, pytensor.function([], value, updates={s: next_c}, mode=mode) + + +def test_host_and_backend_stay_coherent(counter_backend): + s, fn = _counter_function(counter_backend) + + # A host set_value is seen by the compiled read. + s.set_value(5) + assert fn() == 5 # reads 5, advances the native copy to 6 + + # The compiled update is seen by host get_value. + assert s.get_value() == 6 + + # The shared stream continues across interleaved calls, not restarted. + assert fn() == 6 + assert fn() == 7 + + # A later host set_value overrides the backend stream. + s.set_value(0) + assert fn() == 0 # reads 0, advances the native copy to 1 + + # A copy taken while the backend is ahead captures the advanced value, not the + # stale host snapshot (regression for the deepcopy reconcile), and is a fully + # independent shared afterwards -- the two diverge, neither aliases the other. + s_copy = copy.deepcopy(s) + assert s_copy.get_value() == 1 + assert fn() == 1 # advance the original via the backend... + assert s.get_value() == 2 + assert s_copy.get_value() == 1 # ...the copy is untouched + s_copy.set_value(50) # mutate the copy on the host... + assert s_copy.get_value() == 50 + assert s.get_value() == 2 # ...the original is untouched + + +def test_lossy_reconcile_warns_once(lossy_backend): + s, fn = _counter_function(lossy_backend) + fn() # backend advances; host now stale + with pytest.warns(UserWarning, match="lossy"): + s.get_value() + + # Once warned for a source, a later stale read from it does not warn again. + fn() + with warnings.catch_warnings(): + warnings.simplefilter("error") + s.get_value() + + +def test_no_state_without_backend(): + # An ordinary shared no backend ever touches allocates no coherence state. + s = shared(np.array(1.0)) + assert s.container._backend_state is None + s.get_value() + s.set_value(np.array(2.0)) + assert s.container._backend_state is None diff --git a/tests/tensor/random/test_type.py b/tests/tensor/random/test_type.py index 5ff7037c71..142c34bf68 100644 --- a/tests/tensor/random/test_type.py +++ b/tests/tensor/random/test_type.py @@ -52,16 +52,9 @@ def test_filter(self): with pytest.raises(TypeError): rng_type.filter(1) + # Only Generator instances are valid values; a state dict is not. rng_dict = rng.bit_generator.state - assert rng_type.is_valid_value(rng_dict) is False - assert rng_type.is_valid_value(rng_dict, strict=False) - - rng_dict["state"] = {} - - assert rng_type.is_valid_value(rng_dict, strict=False) is False - - rng_dict = {} assert rng_type.is_valid_value(rng_dict, strict=False) is False def test_values_eq(self): @@ -88,13 +81,9 @@ def test_values_eq(self): assert rng_type.values_eq(bitgen_g, bitgen_h) assert rng_type.is_valid_value(bitgen_a, strict=True) - assert rng_type.is_valid_value(bitgen_b.bit_generator.state, strict=False) assert rng_type.is_valid_value(bitgen_c, strict=True) - assert rng_type.is_valid_value(bitgen_d.bit_generator.state, strict=False) assert rng_type.is_valid_value(bitgen_e, strict=True) - assert rng_type.is_valid_value(bitgen_f.bit_generator.state, strict=False) assert rng_type.is_valid_value(bitgen_g, strict=True) - assert rng_type.is_valid_value(bitgen_h.bit_generator.state, strict=False) def test_may_share_memory(self): bg_a = np.random.PCG64()