.WIP new explicit RNG API

Author: ricardoV94

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/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 = False,
+        **kwargs,
+    ):
         if dtype is None:
             dtype = self.dtype
         if dtype == "floatX":
@@ -332,15 +341,26 @@ 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:
+            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/var.py

@@ -4,9 +4,10 @@
 
 from pytensor.compile.sharedvalue import SharedVariable, shared_constructor
 from pytensor.tensor.random.type import random_generator_type
+from pytensor.tensor.random.variable import RandomGeneratorVariable
 
 
-class RandomGeneratorSharedVariable(SharedVariable):
+class RandomGeneratorSharedVariable(SharedVariable, RandomGeneratorVariable):
     def __str__(self):
         return self.name or f"RNG({self.container!r})"
 

pytensor/tensor/random/variable.py

@@ -0,0 +1,59 @@
+import warnings
+from functools import wraps
+from typing import TypeAlias
+
+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 hasattr(self.tag, "used"):
+            warnings.warn(
+                f"RandomGeneratorVariable {self} has already been used. "
+                "You probably want to use the RandomGeneratorVariable that was returned then.",
+                UserWarning,
+            )
+        self.tag.used = True
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
+
+class _random_generator_py_operators:
+    # These can't work because Python requires native output types
+    def __bool__(self):
+        return True
+
+    @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)