What is the status of nnx and TrainState? Are they compatible?

[samuela]

Looking at the nnx docs and example code, I am noticing that TrainState is conspicuously absent. Furthermore, the TrainState docstring still references the outdated linen API.

What's the word on the street on TrainState? Is it still relevant in a flax.nnx world?

[samuela]

To complicate things further, I just discovered https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/helpers.html#flax.nnx.TrainState which is apparently different from https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.train_state.TrainState.

[samuela]

There's also https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.train_state.TrainState but it does not mention flax.nnx.TrainState.

[cgarciae]

Hey @samuela, nnx.TrainState is not documented because it was a prototype from the old days, the need for this API was reduced when we introduced nnx.Optimizer which is more ergonomic. I think we should delete it and instead show users how to use train_state.TrainState with NNX.

[samuela]

I think we should delete it and instead show users how to use train_state.TrainState with NNX.

My read from https://flax.readthedocs.io/en/latest/guides/linen_to_nnx.html#using-trainstate-in-flax-nnx, and specifically "In Flax NNX, [bundling parameters and optimizer] is not really necessary." was that TrainState is now obsolete. But perhaps this incorrect?

[cgarciae]

I don't think that TrainState is obsolete in the sense that there's still many valid use cases for using a functional training loops, this style should still be available to users.