Explicit RNG API

[ricardoV94]

I've been thinking we should cleanup the PyTensor RNG API to be more transparent:

import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.tensor.random.utils import RandomStream

# Current "recommended" API
rng = RandomStream(0)
x = rng.normal()
y = rng.normal(y)
fn = pytensor.function([], y)
print(fn(), fn())  # different numbers due to default updates

# Proposal for "explicit" API that deprecates RandomStream
first_rng = rng = pt.random.rng("rng")
rng, x = rng.normal()
rng, y = rng.normal(x)
# If you don't set mutable=True, nothing terrible happens, 
# just the first use requires a copy
fn = pytensor.function([pytensor.In(first_rng, mutable=True)], [rng, y])
rng_val = np.random.default_rng(0)
rng_val, y1 = fn(rng_val)
rng_val, y2 = fn(rng_val)
print(y1, y2)

# Shared variables can still be used with updates, but rng is still explicit during graph construction
first_rng = rng = pytensor.shared(np.random.default_rng(0), name="rng")
rng, x = rng.normal()
rng, y = rng.normal(x)
# Updates already imply mutable=True
fn = pytensor.function([], y, updates={first_rng: rng})
print(fn(), fn())

# Reusing RNG warns just to avoid obvious mistakes
# UserWarning: RandomGeneratorVariable rng has already been used. You probably want to use the RandomGeneratorVariable that was returned then.
#   warnings.warn(
first_rng.normal()

It looks almost like numpy, except you get the next rng as an output, and you must use that afterwards. I added a warning if you try to call a method from the same rng twice.

It's not unlike the pattern x = x[idx].set(y), where you have to use the next x if you want it to be the version with y in it. All coming from the same requirement of immutability.

This would replace the RandomStream helper, and more importantly the whole SharedVariable.default_update hack that exists solely to try and hide the next_rng thing, and get it picked up by function. This also adds some thorny logic inside Scan, which tries to figure out if there was a new rng.foo() call inside the step function, and get default updates.

All this is quite problematic because you can do graph cloning/manipulation correctly as you would need to check all shared variables default_update and decide what to do with those. Related to #1704

May be worth reviewing implementation details in https://pytensor.readthedocs.io/en/latest/tutorial/prng.html#prng

[Armavica]

I really like it!

[ricardoV94]

Dirty draft PR: #1707

[theorashid]

I like the way jax does explicit rng. Anything closer that (like this proposal) I prefer

[ricardoV94]

JAX uses splitting but that doesn't work well for numpy generators which are meant to be used sequentially. Discussion here: data-apis/array-api#431

The explicit functional part is more similar I agree, but more footgunny, which is why I felt the need for that warning

[cetagostini]

Like it