Update state using pure_callback(...)

[AakashKumarNain]

I have a seed generator defined in a py file named as utils.py. The seed generator is defined like this:

class SeedGenerator:
    def __init__(self, key):
        self.key = key
    
    @jax.jit
    def get_next_key(self):
        key, sub_key = jax.random.split(self.key)
        self.key = key
        return self

    def _tree_flatten(self):
        children = (self.key,)
        aux_data = {}
        return (children, aux_data)

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

# Register the class as a pytree
tree_util.register_pytree_node(SeedGenerator,
                           SeedGenerator._tree_flatten,
                           SeedGenerator._tree_unflatten)

I want to import this seed generator in another py file and update the seed from it. For example,:

from utils import SeedGenerator

def add_random_noise(state, inputs, noise_shape):
    noise = jax.random.normal(state.key, shape=noise_shape)
    return inputs * noise

class AddNoise:
    def __init__(self, noise_shape=None):
        self.seed = SeedGenerator(jax.random.PRNGKey(42))
        self.noise_shape = noise_shape
    
    
    @jax.jit
    def __call__(self, inputs):
        if self.noise_shape is None:
            noise_shape = inputs.shape
            out = add_random_noise(self.seed, inputs, noise_shape)
        else:
            out = add_random_noise(sel.seed, inputs, self.noise_shape)
        
        ## Update the seed with no side-effects here

        return out
    
    def _tree_flatten(self):
        children = (self.seed,)
        aux_data = {"noise_shape": self.noise_shape}
        return (children, aux_data)

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

Is there any easy way to do this?

[jakevdp]

The "pure" in pure_callback refers to the fact that it is built for use with pure functions, i.e functions without side-effects. Updating an internal state is a side-effect, so pure_callback does not have the semantics you are looking for.

If you're looking for callbacks that expect such side-effects, you might try jax.experimental.io_callback.

That said, even with io_callback I'm not sure it's possible to make your code work as written. In particular here:

        ## Update the seed with no side-effects here

This is happening within a JIT-compiled function that does not return self, so there is no way to update self.seed without it being a side-effect.

Some resources for further reading:

• External Callbacks In JAX has an example of io_callback which is quite similar to what you have in mind (updating an external random state)
• JAX Sharp Bits: Pure Functions talks about pure and impure functions in JAX

One last comment: overall, this looks like it may be an attempt to side-step the functional design of JAX PRNGs, and my advice would be: rather than trying to get around it, embrace it 😁.

[AakashKumarNain]

Thanks for the detailed response @jakevdp

This is happening within a JIT-compiled function that does not return self, so there is no way to update self.seed without it being a side-effect.

Yes, I am aware of it. That is why I didn't even write the code block :)

One last comment: overall, this looks like it may be an attempt to side-step the functional design of JAX PRNGs, and my advice would be: rather than trying to get around it, embrace it 😁.

PRNG design is one of the things that I absolutely love about JAX. I just wanted to avoid writing seed generator in the same file. If I use it in the same file, I don't even have to use it within the class, as I can easily mutate the seed outside of it, and pass it again.

Anyway, thanks a lot for the external callback links