Use Numba's new `RandomState` and `Generator` support
As an effort towards #801 , We want to implement a new class based RandomState in our Numba backend. It'll allow us to curcumvent Numba's limitation of a global random state and will give the ability to have multiple independent random states just like Numpy's new RandomState interface.
The current implementation for this is at: https://gist.github.com/kc611/821b41413d0f6cc0ad34d842a25aa7ea
The general idea is to synchronize Numba's global random state with the current RandomStateNumbaType's state key before each random draw is made. The new global random state can then be put back into the RandomStateNumbaType object, allowing local state based draws.
If we want to go a step further we might actually want to implement a complete RandomState Numba equivalent and that might be actually easier since we'd just have to move the logic in
https://github.com/numba/numba/blob/39f03307113eaf76f0b1308cfbe0c8bf78a5c3e9/numba/cpython/randomimpl.py#L115-L241
In the new Type class. This is something we'd eventually want to back-port to Numba but we should at-least have a proper working implementation before that.
If we want to go a step further we might actually want to implement a complete
RandomStateNumba equivalent and that might be actually easier since we'd just have to move the logic inhttps://github.com/numba/numba/blob/39f03307113eaf76f0b1308cfbe0c8bf78a5c3e9/numba/cpython/randomimpl.py#L115-L241
In the new
Typeclass. This is something we'd eventually want to back-port to Numba but we should at-least have a proper working implementation before that.
Yes, that sounds ideal; especially if it helps pave the way for Numba support for the new Generator types.
The important parts of this are being implemented here: https://github.com/numba/numba/projects/19#card-84051388.
Update: The Numba changes are mostly done (and released). One of the last things we need is the ability to lower NumPyRandomGeneratorTypes (see https://github.com/numba/numba/issues/4499#issuecomment-1229235923).
Let me take a look at this once I'm done with my open PRs. That and Scan are in a way of being able to compile Aesara models and samplers to Numba.