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.

aesara/aesara/link/numba/dispatch/random.py

Lines 51 to 79 in a0604e5

	@unbox(RandomStateNumbaType)
	def unbox_random_state(typ, obj, c):
	"""Convert a `RandomState` object to a native `RandomStateNumbaModel` structure.
	Note that this will create a 'fake' structure which will just get the
	`RandomState` objects accepted in Numba functions but the actual information
	of the Numba's random state is stored internally and can be accessed
	anytime using ``numba._helperlib.rnd_get_np_state_ptr()``.
	"""
	interval = cgutils.create_struct_proxy(typ)(c.context, c.builder)
	is_error = cgutils.is_not_null(c.builder, c.pyapi.err_occurred())
	return NativeValue(interval._getvalue(), is_error=is_error)
	@box(RandomStateNumbaType)
	def box_random_state(typ, val, c):
	"""Convert a native `RandomStateNumbaModel` structure to an `RandomState` object
	using Numba's internal state array.
	Note that `RandomStateNumbaModel` is just a placeholder structure with no
	inherent information about Numba internal random state, all that information
	is instead retrieved from Numba using ``_helperlib.rnd_get_state()`` and a new
	`RandomState` is constructed using the Numba's current internal state.
	"""
	pos, state_list = _helperlib.rnd_get_state(_helperlib.rnd_get_np_state_ptr())
	rng = RandomState()
	rng.set_state(("MT19937", state_list, pos))
	class_obj = c.pyapi.unserialize(c.pyapi.serialize_object(rng))
	return class_obj

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: numba/numba#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 #662 (i.e. lack of Generator support in Numba), and the work involved is covered by #814.