Is this project still maintained?

[patrick-kidger]

Hey @paganpasta! I often point people at Eqxvision, as it's a great model zoo. (And one that you've clearly put a lot of work into.)

I was curious if this project was still maintained? I notice things don't seem to have been updated to eqx.nn.BatchNorm (which has been stable for some months now).

If it is maintained, I'd like to see if we could promote this project a little more. But if it's not, then no worries -- I understand maintenance is a burden, and there's libraries I've previously decided to stop supporting as well.

[paganpasta]

Hi,

I didn't get a chance to bring it up to speed in the past months. I do intend to update the project and as you pointed out things are stale prior to BatchNorm.

Would it be possible for you to point me in the direction of changes (before/after) required for BatchNorm integration now ? I suspect this will be the largest of the changes required to bring the repo up to speed.

[patrick-kidger]

Ah, marvellous! That's great to hear.

The main change is that the statefulness of BatchNorm is now explicitly threaded through: instead of x = norm(x), we have x, state = norm(x, state). This gives us an explicit place to store the updated means/variances, without weird side-effecting hackery. This means we're now compatible with arbitrary JAX transformations, rather than just kind of exploding when placed inside a pmap or checkpoint or scan!

This does manifest as an API change, of course. Statefulness is something that the surrounding module has to be aware of, so that it can thread the state in and out.

One convenience worth knowing: LayerNorm and GroupNorm deliberately expose the same API (if passed a state, they will just return it unchanged), so that all three layers are easily interchangeable.

Here's a full example for what it looks like to use a stateful module: https://docs.kidger.site/equinox/examples/stateful/

[paganpasta]

Thanks! I'll get started on it.