Update a random number generator in a jitted function without tracing it

[gianlucadetommaso]

I'm trying to update a random number generator in a jit function without tracing it. Here is an example code:

from jax import random, pure_callback

class Foo:
    rng = random.PRNGKey(0)

    @partial(jax.jit, static_argnums=(0,))
    def fun(self):
        def _update():
            self.rng = random.split(self.rng)[0]
        jax.pure_callback(_update, None)

foo = Foo()
foo.fun()
print(foo.rng)

This gives DeviceArray([0, 0], dtype=uint32).

If I comment out the decorator, I obtain DeviceArray([4146024105, 967050713], dtype=uint32), which is the desired result.

Any idea how I can obtain DeviceArray([4146024105, 967050713], dtype=uint32) (i.e. updated rng) also if the function is jitted?

[jakevdp]

The problem is that you've marked self as static, but it does not have a meaningful hash value. See FAQ: how to use JIT with methods for a description of the issue and some potential solutions.

[pkairys]

I'm running into. similar issue. Maybe I am misunderstanding something?

class CustomClass:
    def __init__(self,key):
        self.key = key
    @jit
    def newkey(self):
        newkey, subkey = random.split(self.key)
        self.key = newkey

    def _tree_flatten(self):
        children = (self.key,)  # arrays / dynamic values
        aux_data = {}  # static values
        return (children, aux_data)

    @classmethod
    def _tree_unflatten(cls, aux_data, children):
        return cls(*children)

from jax import tree_util
tree_util.register_pytree_node(CustomClass,
                           CustomClass._tree_flatten,
                           CustomClass._tree_unflatten)

If I remove the @jit method then the key is updated appropriately. However, with @jit the key is not updated, even though I've marked the attribute as dynamic. Perhaps I'm doing something wrong?

[jakevdp]

You're running into this sharp bit: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions

JAX's jit and other transformations assume your functions are pure – in particular, that the function has no side-effects. Here you are updating self.key as a side-effect. Instead you might do something like this:

class CustomClass:
    def __init__(self,key):
        self.key = key
    @jit
    def with_newkey(self):
        newkey, subkey = random.split(self.key)
        self.key = newkey
        return self

    def _tree_flatten(self):
        children = (self.key,)  # arrays / dynamic values
        aux_data = {}  # static values
        return (children, aux_data)

    @classmethod
    def _tree_unflatten(cls, aux_data, children):
        return cls(*children)

from jax import tree_util
tree_util.register_pytree_node(CustomClass,
                           CustomClass._tree_flatten,
                           CustomClass._tree_unflatten)

obj = CustomClass(random.PRNGKey(0))
print(obj.key)  # [0 0]

obj = obj.with_newkey()  # note: update to `obj` must be explicit!
print(obj.key)  # [4146024105  967050713]

[pkairys]

I see, thank you for the additional information!

I incorrectly assumed that since I had marked self.key as a dynamic variable, I'd be able to change it because it seems like the object self is getting passed into the scope of the method. But I guess that isn't quite true for the methods.

Either way, thanks for the clarification!