aesara icon indicating copy to clipboard operation
aesara copied to clipboard

Use Numba's new `RandomState` and `Generator` support

Open kc611 opened this issue 3 years ago • 5 comments

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.

kc611 avatar Feb 09 '22 16:02 kc611

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.

kc611 avatar Feb 09 '22 17:02 kc611

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.

Yes, that sounds ideal; especially if it helps pave the way for Numba support for the new Generator types.

brandonwillard avatar Feb 09 '22 19:02 brandonwillard

The important parts of this are being implemented here: https://github.com/numba/numba/projects/19#card-84051388.

brandonwillard avatar Aug 04 '22 16:08 brandonwillard

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).

brandonwillard avatar Sep 21 '22 21:09 brandonwillard

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.

rlouf avatar Sep 21 '22 21:09 rlouf