Avital Oliver

If anyone has gotten ImageNet training to work, I'd be thrilled to get some pointers for getting started. So far I've been getting: ``` Epoch: [0][0/7317] Time 62.236 (62.236) Data...

IIUC with your new proposal, you must use `preserve_identities(jax.jit)` rather than `jax.jit`. And the same is true for all other JAX transformations. But what if you call into another library...

Oh! Curious what the edge cases are? (And I'm surprised by the code snippet you posted last because it doesn't set `self.counter2 = self.counter1`)

> I understand. But the issue with that is that you open users to bugs by allowing impure methods. The reality is that Jax's decorated functions (`jit`, `grad`, etc.) are...

> > Typically the code in the main training loop isn't pure anyways (and isn't meant to be pure, as it reports metrics, saves checkpoints, etc). > Could you point...

The problem with any hoop isn't it's complexity -- it's that it's something you have to learn suddenly, when you "just wanted to try this one thing". So any hoop...

I guess another way to put it -- if someone *really* wants immutable data structures, they can always do, e.g. `FrozenDict(module.init(...))`. So the question is: which default serves the users...

I agree, those are useful. I don't know what the long term best place for them is, though. Maybe even outside of Flax? So I guess maybe it's fine to...

Just curious @n2cholas -- are you planning to work on this or did you file it for us to keep track of it (both answers are fine, just checking). If...

As a first step, yes. Add a deprecation warning and make sure it is clearly deprecated in the docs. Make sure we don't use them anywhere in our examples/HOWTOs/etc. Then...