Implementing Batch Normalization

In Flax, Batch Normalization is a bit finicky since each call to apply requires marking batch_stats as mutable and updating the batch_stats afterward.

bn = flax.linen.BatchNorm(use_running_average=True)

x = jnp.arange(24).reshape(3, 6)

vars = bn.init(random.PRNGKey(0), x)

# Mark the batch stats as mutable so we can update them in the variable dictionary
x_normed, mutated_vars = bn.apply(vars, x, mutable=['batch_stats'])

vars = {**vars, **mutated_vars}  # Update the variables with our diff

x_normed2, mutated_vars2 = bn.apply(vars, x, mutable=['batch_stats'])

How could this be implemented as a Module in Equinox? I’m happy to submit an implementation given some guidance.

[patrick-kidger]

Yeah, this is a known issue with JAX – namely, that host_callback.call handles errors differently depending on OS, device, or phase of the moon.

See also https://github.com/google/jax/issues/9457

I don’t think there’s much that can be done about this one from the point of view of Equinox I’m afraid.

[patrick-kidger]

Well, I got successfully nerd-sniped into spending my weekend implementing this. (Mostly the new “stateful” technology that makes this possible.)

equinox.experimental.BatchNorm now exists. Happy hacking.