keeping a state

[kyunghyuncho]

hi,

a jax newbie here. thanks for creating a great framework!

i'm trying to figure out what's the best way to keep some running stats of computation across jax functions. for instance, i'm trying to implement batch normalization as a class, and this class keeps the running averages of the mean and variance of the batch statistics. see, e.g., the code snippet below.

    def forward(self, p, b, x):
        x_ = jnp.transpose(x, [0, 2, 3, 1]).reshape(-1, x.shape[1])
        mu = x_.mean(0)
        var = x_.var(0)
        x_ = (x_ - mu)/(jnp.sqrt(var) + 1e-6) * p['gamma'] + p['alpha']
        x_ = x_.reshape(x.shape[0], x.shape[2], x.shape[3], -1)
        x_ = jnp.transpose(x_, [0, 3, 1, 2])
        
        b['mu'] = self.coeff * b['mu'] + (1.-self.coeff) * mu
        b['var'] = self.coeff * b['var'] + (1.-self.coeff) * var
        
        return x_

p is a dictionary containing the parameters (alpha and gamma), and b is a buffer containing the running statistics (mu and var). this of course raises an error by the tracer, as the buffer's content is stateful.

i understand this goes against the jax's implementation and philosophy and that i need to think of a different way (e.g., return the new buffer content together with the output of the computation,) but there's a huge inertia for me as the current pytorch user to want to keep some states across computation.

is there any other way around than to put everything in as a part of the arguments and return everything out as a part of the output to keep any state? or, perhaps, is there any plan to introduce some kind of global tensor storage for jax that can be easily accessible from any jax-based functions? i totally understand this goes against the philosophy, but i was just thinking aloud my wish.

cheers,
-- k

[shoyer]

JAX doesn't (yet) have any builtin support for mutable state, but this is something you can find in a number of higher level neural net libraries build on top of JAX.

For two examples, see haiku.transform_with_state and flax.nn.stateful.

[kyunghyuncho]

thanks for the pointer. i also ended up implementing each forward function to optionally return an updated buffer state (based on the buffer state passed to the function as an argument.) it was much more straightforward than i initially thought, but because almost any framework that wraps jax would need their own way to implement stateful computation, it would be worth an effort in the future to support it natively at the level of jax.

thanks again!

• k

[mattjj]

We've been thinking hard about this one for a while, with collaborators from the Flax, Haiku, Trax, and Oryx teams too. It's still a work in progress though.

Basically, +1 to what @shoyer said, with the extra emphasis that we're interested in doing better here.