Use Numba Generators for random graphs and deprecate shared RandomState variables

[ricardoV94]

Description

RandomState are legacy in numpy and we can save some complexity by letting go of them in PyTensor.

We were "obliged" to keep them because that's the only kind of RNG that Numba supported until now

[jessegrabowski]

I'm trying to understand this issue to start a PR. What actually needs to be done? Numba functions can take numpy random generators without any hassle now. For example, this works:

import numba as nb
import numpy as np

@nb.njit
def draw_nb(rng, loc, scale, size):
    return rng.normal(loc=loc, scale=scale, size=size)

rng = np.random.default_rng()
draw_nb(rng, 0.0, 1.0, (10,))

So would it be enough to just chop out all of the extra machinery from numba\dispatch\random.py related to the random states here? Plus the check in make_numba_random_fn? Or is there something deeper going on.

Numba still doesn't support broadcasting from parameters, so all the rest of the machinery seems like it needs to say (though I personally find it quite difficult to follow, it would be nice to refactor it to be more clear).

[ricardoV94]

Numba still doesn't support broadcasting from parameters, so all the rest of the machinery seems like it needs to say (though I personally find it quite difficult to follow, it would be nice to refactor it to be more clear).

This is the big hurdle. Generators without broadcasting is pretty useless so we have to support it if we want to really support RVs in numba backend (and phase out the RandomState)

[ricardoV94]

This is the PR where Generator support was added to Aesara. The main complexity is writing broadcasting logic with python strings which is the usual numba backend PITA: aesara-devs/aesara#1245

[jessegrabowski]

But I'm saying I'm pretty sure you can directly plug a generator into what they already have? There's even a numba datatype for numpy random generators (nb.types.NumPyRandomGenerator)

[ricardoV94]

I think the only challenge is the broadcasting logic, which I think can't be written as a Python function without writing it for every Op?

I don't remember exactly where did things break.

Also the strict API for RVs requires copying the RNG if the Op is not inplace. Not sure if this is relevant.

[ricardoV94]

What did you mean, what they already have?

[ricardoV94]

Feel free to open a PR if you it seems like some minimal changes do the job (or even if they don't).

Unfortunately, I've lost the context to this issue to be able to help just from thinking

[jessegrabowski]

The broadcasting is done with a loop, so it's actually not too bad. Here is a basic sketch:

import numba as nb
import numpy as np

@nb.njit
def draw_nb(rng, loc, scale, size):
    loc_bcast = np.broadcast_to(loc, size)
    scale_bcast = np.broadcast_to(scale, size)

    bcast_samples = np.empty(size, dtype=np.float64)
    
    for idx in np.ndindex(size):
        bcast_samples[idx] = rng.normal(loc_bcast[idx], scale_bcast[idx])
    return bcast_samples

rng = np.random.default_rng(1)
loc = np.zeros((5, 5))
scale = 1.0
size = (10, 5, 5)
samples_np = rng.normal(loc=loc, scale=scale, size=size)

rng = np.random.default_rng(1)
samples_nb = draw_nb(rng, loc, scale, size)

np.allclose(samples_np, samples_nb) #True

I guess it just seems like all the special packing/unpacking of the random state that is done in the numba linker can just go. But is that the only thing causing problems?

Sure I'll open a PR. I'm also curious why this all causes problems only with certain models (like scan) but not in others.

[ricardoV94]

Can you write that function in a way that works for any RV Op regardless of the number of parameters and ndim_supp/ndims_params?

I imagine that's the abstraction that complicates things.

[ricardoV94]

The boxing and unboxing of RandomState is supposed to go out completely yes

[jessegrabowski]

No, you have to write overloads on a case-by-case basis. But that should be a hassle, not a blocker. Also it already exists.

The current code jumps through a lot of hoops to make a function that generates generalized boilerplate code. I guess it's elegant, but it makes it really hard to read. A bunch of individual functions would be easier to maintain IMO.