The Numba backend does not support `RandomStream`

[rlouf]

The following code uses a RandomStream to define random variables and buid a simple logprob function with joint_logprob:

from aesara.graph.fg import FunctionGraph
from aesara.link.numba.dispatch import numba_funcify

import aesara.tensor as at

from aeppl import joint_logprob

srng = at.random.RandomStream(seed=0)
Y_rv = srng.normal(1, 2)

def logprob_fn(y):
    logprob = joint_logprob({Y_rv: y})
    return logprob

yrv_fgraph = FunctionGraph(outputs=[Y_rv], clone=False)
fn = numba_funcify(yrv_fgraph)

  File "<stdin>", line 20, in <module>
  File "<stdin>", line 18, in main
  File "/usr/lib/python3.9/functools.py", line 877, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/home/remi/projects/aesara/aesara/link/numba/dispatch/basic.py", line 354, in numba_funcify_FunctionGraph
    return fgraph_to_python(
  File "/home/remi/projects/aesara/aesara/link/utils.py", line 723, in fgraph_to_python
    compiled_func = op_conversion_fn(
  File "/usr/lib/python3.9/functools.py", line 877, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/home/remi/projects/aesara/aesara/link/numba/dispatch/random.py", line 221, in numba_funcify_RandomVariable
    raise TypeError("Numba does not support NumPy `Generator`s")
TypeError: Numba does not support NumPy `Generator`s

It is problematic since the current implementation of HMC/NUTS relies on RandomStreams and its autoupdate feature inside of Scan ops. Is there any fundamental incompatibility or is it just not implemented yet?

[kc611]

For supporting something like that in Numba backend (calling RandomStream and its method from within Numba code) we'll need to implement boxing/unboxing for the RandomStream class similar to how we have did for RandomState.

https://github.com/aesara-devs/aesara/blob/a0604e5224fbfd4d4bfa804b1c3ac101cdacd52c/aesara/link/numba/dispatch/random.py#L51-L79

The main problem however is that Numba still relies on the old seeding interface i.e. It only has global seeding and not class based seeding logic like the one relied on by RandomStream, so that can be a problem.

[kc611]

There's an open issue for this in Numba: https://github.com/numba/numba/issues/4499

[rlouf]

Now that Generators are supported in Numba, we can go ahead and support RandomStream in the Numba backend.

[brandonwillard]

This is effectively a duplicate of https://github.com/aesara-devs/aesara/issues/662 (i.e. lack of Generator support in Numba), and the work involved is covered by https://github.com/aesara-devs/aesara/issues/814.