Updates on flax, jax and optax caused regressions in google-research projects

[xvdp]

Who is responsible for maintenance of deprecated code impact on other parts of your org?
Project https://github.com/google-research/google-research/tree/master/diffusion_distillation, a very nice paper from 2022, used flax.optim and a custom TrainState based on @flax.struct.dataclass which have since been replaced by optax and flax.training.train_state.TrainState. Im sure it is great but between google colab forcing numpy, jax, flax etc.. versions, as well as python 3.12 as well as jaxlib 0.4.1 no longer being in pypi and custom installation from wheel containg no cuda versions... it is impossible to recreate old projects. Especially if one (partly for this reason of deprecations) leans towards torch. * I do prefer jax in principle but I dont use it because "old" projects are broken (3 years old stuff that is very relevatn should not be "old")

I tried to fix the project but I only did it partially, something in either flax.optim -> optax or Trainstate -> flax.training.train_state.TrainState breaks my attempted fix.

google-research/google-research#2976

Can someone quickly fix that? Or suggest what in these 2 files i gotta change. to make a pull request ?
https://github.com/xvdp/google-research/blob/fixflaxversion/diffusion_distillation/diffusion_distillation/model.py
https://github.com/xvdp/google-research/blob/fixflaxversion/diffusion_distillation/diffusion_distillation.ipynb

[xvdp]

downgrading jax and flax to 0.4.1. even if it is possible locally, wont allow to run this project, as jax functions have been updated
xvdp/google-research@53f52d1
so jaxlib 0.4.1, will fail as jax has no tree attribute in the orginal.

I was trying to replicate a single loop of the original code optimizer to see where my update is wrong, but i cant.
Someone pls fix this project. thanx

[vfdev-5]

@xvdp thanks for the report and your opinion. I understand that compatibility issues can be very annoying. I think the 3rd party projects and their maintainers should adapt their code to the new versions of jax, flax etc. We can help with answering questions but unlikely fix their code.
If you would like to help you can follow their contributing guides and submit a PR.

Personally, I would try to run that project with the versions of the publication (sad that they haven't pinned the versions: https://github.com/google-research/google-research/blob/master/diffusion_distillation/diffusion_distillation/requirements.txt)

[xvdp]

thanx. im sure you know about the bane of code rot. And since ~ 2021 ML adjacent rot has hit ludicrous speed. But we are all working with stuff which seminal papers date mostly from 2013-2017, eg. GAN,Flow, Diffusion, VAE, Transformers, Unets, Skip connections, Batch Norm also Distillation and Adversarial. After that it is all derivatives permitted by understanding of flow of gradients, and latents, training, and scaling, and data alignment.
My code may be crap but google-research's could tick along without rot, and in the meanwhile it could be leveraged used to teach gemini+causal logic, the meaning of systems, TDD passing of time and so on.
Whatevs, i fixed the error, but with a bug, still got the gist of the distill step, just a pity it isnt kept up to date.