What is the status of nnx and TrainState? Are they compatible?
Open
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
Looking at the nnx [docs](https://flax.readthedocs.io/en/latest/why.html) and [example code](https://flax.readthedocs.io/en/latest/mnist_tutorial.html), I am noticing that [TrainState](https://flax.readthedocs.io/en/v0.8.3/api_reference/flax.training.html#train-state) 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?
Contributor guide
Assessment
This issue has not been assessed yet.