.WIP new explicit RNG API

Author: ricardoV94

pytensor/__init__.py

@@ -159,7 +159,7 @@ def get_underlying_scalar_constant(v):
 
 
 # isort: off
-import pytensor.tensor.random.var
+import pytensor.tensor.random.variable
 import pytensor.sparse
 from pytensor.ifelse import ifelse
 from pytensor.scan import checkpoints

pytensor/compile/function/__init__.py

@@ -1,6 +1,7 @@
 import logging
 import re
 import traceback as tb
+import warnings
 from collections.abc import Iterable
 from pathlib import Path
 
@@ -102,7 +103,7 @@ def function(
     givens: Iterable[tuple[Variable, Variable]]
     | dict[Variable, Variable]
     | None = None,
-    no_default_updates: bool = False,
+    no_default_updates: bool | None = None,
     accept_inplace: bool = False,
     name: str | None = None,
     rebuild_strict: bool = True,
@@ -266,6 +267,15 @@ def opt_log1p(node):
         of just writing it in C from scratch.
 
     """
+    if no_default_updates is not None:
+        warnings.warn(
+            "The no_default_updates parameter is deprecated and will be "
+            "removed in a future version of PyTensor. Please set updates manually ",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+    else:
+        no_default_updates = False
     if isinstance(outputs, dict):
         assert all(isinstance(k, str) for k in outputs)
 

pytensor/compile/sharedvalue.py

@@ -1,6 +1,7 @@
 """Provide a simple user friendly API to PyTensor-managed memory."""
 
 import copy
+import warnings
 from contextlib import contextmanager
 from functools import singledispatch
 from typing import TYPE_CHECKING
@@ -161,6 +162,7 @@ def default_update(self) -> Variable | None:
 
     @default_update.setter
     def default_update(self, value):
+        warnings.warn("Setting default_update is deprecated.", DeprecationWarning)
         if value is not None:
             self._default_update = self.type.filter_variable(value, allow_convert=True)
         else:

pytensor/tensor/random/__init__.py

@@ -4,3 +4,4 @@
 from pytensor.tensor.random.basic import *
 from pytensor.tensor.random.op import default_rng
 from pytensor.tensor.random.utils import RandomStream
+from pytensor.tensor.random.variable import rng

pytensor/tensor/random/basic.py

@@ -285,7 +285,7 @@ def __call__(self, loc=0.0, scale=1.0, size=None, **kwargs):
 normal = NormalRV()
 
 
-def standard_normal(*, size=None, rng=None, dtype=None):
+def standard_normal(*, size=None, rng=None, dtype=None, **kwargs):
     """Draw samples from a standard normal distribution.
 
     Signature
@@ -302,7 +302,7 @@ def standard_normal(*, size=None, rng=None, dtype=None):
         is returned.
 
     """
-    return normal(0.0, 1.0, size=size, rng=rng, dtype=dtype)
+    return normal(0.0, 1.0, size=size, rng=rng, dtype=dtype, **kwargs)
 
 
 class HalfNormalRV(ScipyRandomVariable):
@@ -516,7 +516,7 @@ def chisquare(df, size=None, **kwargs):
     return gamma(shape=df / 2.0, scale=2.0, size=size, **kwargs)
 
 
-def rayleigh(scale=1.0, *, size=None, **kwargs):
+def rayleigh(scale=1.0, *, size=None, return_next_rng=False, **kwargs):
     r"""Draw samples from a Rayleigh distribution.
 
     The probability density function for `rayleigh` with parameter `scale` is given by:
@@ -550,7 +550,13 @@ def rayleigh(scale=1.0, *, size=None, **kwargs):
     scale = as_tensor_variable(scale)
     if size is None:
         size = scale.shape
-    return sqrt(chisquare(df=2, size=size, **kwargs)) * scale
+    next_rng, chisquare_draws = chisquare(
+        df=2, size=size, return_next_rng=True, **kwargs
+    )
+    rayleigh_draws = sqrt(chisquare_draws) * scale
+    if return_next_rng:
+        return next_rng, rayleigh_draws
+    return rayleigh_draws
 
 
 class ParetoRV(ScipyRandomVariable):
@@ -1986,7 +1992,7 @@ def rng_fn(self, *params):
         return out
 
 
-def choice(a, size=None, replace=True, p=None, rng=None):
+def choice(a, size=None, replace=True, p=None, rng=None, return_next_rng=False):
     r"""Generate a random sample from an array.
 
 
@@ -2016,17 +2022,23 @@ def choice(a, size=None, replace=True, p=None, rng=None):
         # This is equivalent to the numpy implementation:
         # https://github.com/numpy/numpy/blob/2a9b9134270371b43223fc848b753fceab96b4a5/numpy/random/_generator.pyx#L905-L914
         if p is None:
-            idxs = integers(0, a_size, size=size, rng=rng)
+            next_rng, idxs = integers(
+                0, a_size, size=size, rng=rng, return_next_rng=True
+            )
         else:
-            idxs = categorical(p, size=size, rng=rng)
+            next_rng, idxs = categorical(p, size=size, rng=rng, return_next_rng=True)
 
         if a.type.ndim == 0:
             # A was an implicit arange, we don't need to do any indexing
             # TODO: Add rewrite for this optimization if users passed arange
-            return idxs
-
-        # TODO: Can use take(a, idxs, axis) to support numpy axis argument to choice
-        return a[idxs]
+            out = idxs
+        else:
+            # TODO: Can use take(a, idxs, axis) to support numpy axis argument to choice
+            out = a[idxs]
+        if return_next_rng:
+            return next_rng, out
+        else:
+            return out
 
     # Sampling with p is not as trivial
     # It involves some form of rejection sampling or iterative shuffling under the hood.
@@ -2063,7 +2075,7 @@ def choice(a, size=None, replace=True, p=None, rng=None):
     op = ChoiceWithoutReplacement(signature=signature, dtype=dtype)
 
     params = (a, core_shape) if p is None else (a, p, core_shape)
-    return op(*params, size=None, rng=rng)
+    return op(*params, size=None, rng=rng, return_next_rng=return_next_rng)
 
 
 class PermutationRV(RandomVariable):

pytensor/tensor/random/op.py

@@ -314,7 +314,16 @@ def infer_shape(self, fgraph, node, input_shapes):
 
         return [None, list(shape)]
 
-    def __call__(self, *args, size=None, name=None, rng=None, dtype=None, **kwargs):
+    def __call__(
+        self,
+        *args,
+        size=None,
+        name=None,
+        rng=None,
+        dtype=None,
+        return_next_rng: bool | None = None,
+        **kwargs,
+    ):
         if dtype is None:
             dtype = self.dtype
         if dtype == "floatX":
@@ -332,15 +341,31 @@ def __call__(self, *args, size=None, name=None, rng=None, dtype=None, **kwargs):
             props["dtype"] = dtype
             new_op = type(self)(**props)
             return new_op.__call__(
-                *args, size=size, name=name, rng=rng, dtype=dtype, **kwargs
+                *args,
+                size=size,
+                name=name,
+                rng=rng,
+                dtype=dtype,
+                return_next_rng=return_next_rng,
+                **kwargs,
             )
 
-        res = super().__call__(rng, size, *args, **kwargs)
-
+        node = self.make_node(rng, size, *args)
+        outputs = node.outputs
         if name is not None:
-            res.name = name
-
-        return res
+            outputs[self.default_output].name = name
+        if return_next_rng:
+            return outputs
+        else:
+            if return_next_rng is None:
+                warnings.warn(
+                    "The default behavior of RandomVariable.__call__ is changing to return both the next RNG and the draws. "
+                    "Please set return_next_rng explicitly to avoid this warning.",
+                )
+            out = outputs[self.default_output]
+            if kwargs.get("return_list", False):
+                return [out]
+            return out
 
     def make_node(self, rng, size, *dist_params):
         """Create a random variable node.

pytensor/tensor/random/utils.py

@@ -1,3 +1,4 @@
+import warnings
 from collections.abc import Callable, Sequence
 from functools import wraps
 from itertools import zip_longest
@@ -231,6 +232,11 @@ def __init__(
             [np.random.SeedSequence], np.random.Generator
         ] = np.random.default_rng,
     ):
+        warnings.warn(
+            "pytensor.tensor.random.utils.RandomStream is deprecated and will be removed in a future release.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
         if namespace is None:
             from pytensor.tensor.random import basic  # pylint: disable=import-self
 

pytensor/tensor/random/var.py

@@ -0,0 +1,91 @@
+import copy
+import warnings
+from functools import wraps
+from typing import TypeAlias
+
+import numpy as np
+
+from pytensor import config
+from pytensor.compile.sharedvalue import SharedVariable, shared_constructor
+from pytensor.graph.basic import OptionalApplyType, Variable
+from pytensor.tensor.random.basic import normal
+from pytensor.tensor.random.type import RandomGeneratorType, random_generator_type
+from pytensor.tensor.variable import TensorVariable
+
+
+RNG_AND_DRAW: TypeAlias = tuple["RandomGeneratorVariable", TensorVariable]
+
+
+def warn_reuse(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if getattr(self.tag, "used", False) and config.warn_rng.reuse:
+            warnings.warn(
+                f"RandomGeneratorVariable {self} has already been used. "
+                "You probably want to use the new RandomGeneratorVariable that was returned when you used it.",
+                UserWarning,
+            )
+        self.tag.used = True
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
+
+class _random_generator_py_operators:
+    @warn_reuse
+    def normal(self, loc=0, scale=1, size=None) -> RNG_AND_DRAW:
+        return normal(loc, scale, size=size, rng=self, return_next_rng=True)
+
+
+class RandomGeneratorVariable(
+    _random_generator_py_operators,
+    Variable[RandomGeneratorType, OptionalApplyType],
+):
+    """The Variable type used for random number generator states."""
+
+
+RandomGeneratorType.variable_type = RandomGeneratorVariable
+
+
+def rng(name=None) -> RandomGeneratorVariable:
+    """Create a new default random number generator variable.
+
+    Returns
+    -------
+    RandomGeneratorVariable
+        A new random number generator variable initialized with the default
+        numpy random generator.
+    """
+
+    return random_generator_type(name=name)
+
+
+class RandomGeneratorSharedVariable(SharedVariable, RandomGeneratorVariable):
+    def __str__(self):
+        return self.name or f"RNG({self.container!r})"
+
+
+@shared_constructor.register(np.random.RandomState)
+@shared_constructor.register(np.random.Generator)
+def randomgen_constructor(
+    value, name=None, strict=False, allow_downcast=None, borrow=False
+):
+    r"""`SharedVariable` constructor for NumPy's `Generator` and/or `RandomState`."""
+    if isinstance(value, np.random.RandomState):
+        raise TypeError(
+            "`np.RandomState` is no longer supported in PyTensor. Use `np.random.Generator` instead."
+        )
+
+    rng_sv_type = RandomGeneratorSharedVariable
+    rng_type = random_generator_type
+
+    if not borrow:
+        value = copy.deepcopy(value)
+
+    return rng_sv_type(
+        type=rng_type,
+        value=value,
+        strict=strict,
+        allow_downcast=allow_downcast,
+        name=name,
+    )

pytensor/tensor/random/variable.py

@@ -82,7 +82,7 @@ def test_rng_copy():
     rng.type.values_eq(rng.get_value(), np.random.default_rng(123))
 
 
-def test_rng_non_default_update():
+def test_rng_custom_update():
     rng = shared(np.random.default_rng(1))
     rng_new = shared(np.random.default_rng(2))