Boxing/Unboxing of Numpy RandomStates

Numba Support: How do I do ...?
1

I’ve been trying to implement unboxing (and boxing) of Numpy RandomState objects, trying to convert them into Native objects to be able to call its methods from within Numba-Jitted code.

So far this is what I have:

from numba import types
from numba.core import cgutils
from numba.extending import NativeValue, box, models, register_model, typeof_impl, unbox
from numpy.random import RandomState


class RandomStateNumbaType(types.Type):
    def __init__(self):
        super(RandomStateNumbaType, self).__init__(name="RandomState")


random_state_numba_type = RandomStateNumbaType()


@typeof_impl.register(RandomState)
def typeof_index(val, c):
    return random_state_numba_type


@register_model(RandomStateNumbaType)
class RandomStateNumbaModel(models.StructModel):
    def __init__(self, dmm, fe_type):
        members = [
            ("state_key", types.Array(types.uint32, 1, "C")),
        ]
        models.StructModel.__init__(self, dmm, fe_type, members)


@unbox(RandomStateNumbaType)
def unbox_random_state(typ, obj, c):
    """Convert a `RandomState` object to a native `RandomStateNumbaModel` structure."""
    numba_randomstate = cgutils.create_struct_proxy(typ)(c.context, c.builder)
    bitgen = c.pyapi.object_getattr_string(obj, "_bit_generator")
    state_dict = c.pyapi.object_getattr_string(bitgen, "state")
    # TODO: Get the state keys from the `state_dict` Python dictionary object

    # TODO: Store the keys in numba_randomstate as native array

    is_error = cgutils.is_not_null(c.builder, c.pyapi.err_occurred())
    return NativeValue(numba_randomstate._getvalue(), is_error=is_error)


@box(RandomStateNumbaType)
def box_random_state(typ, val, c):
    """Convert a native `RandomStateNumbaModel` structure to an `RandomState` object."""
    rng = RandomState()
    # TODO: Use `val` to get the RandomState keys and put it in rng using rng.set_state

    class_obj = c.pyapi.unserialize(c.pyapi.serialize_object(rng))
    return class_obj

But I cannot figure out the logic / code that’ll go in regions marked by TODOs since that part of code is not in Numba docs. Can somebody point me to the right location for documentation / implementations of such?

Also, The overall plan for this implementation is to have class based random seeds like Numpy does with its new API, within Numba code. I plan to overload/implement methods for RandomStateNumbaType and synchronize Numba’s global state key with the key of the current RandomStateNumbaType every time a random method is called. The updated key can then be put back into the RandomStateNumbaType. Will something like this work?