Dependency modernization: deprecated JAX APIs + future Flax NNX migration

[exopoiesis]

Summary

We've been using CrystalFormer (v0.6) for generative crystal design (Fe-Ni-S sulfide catalysts, ~5000+ structures generated). Works great with JAX 0.4.35 + dm-haiku 0.0.15.

The Haiku question

dm-haiku has been archived by Google DeepMind, with the recommended migration path being Flax NNX. Haiku 0.0.15 works with JAX up to ~0.4.x, but is incompatible with JAX >= 0.5 (named_shape removed from ShapedArray).

Are there any plans to migrate to Flax NNX (or another framework)? If this is something you'd be open to, we'd be happy to work on it in a fork first and submit a PR once validated.

Context

We're using CrystalFormer for Fe-Ni-S mineral generation and recently contributed to JDFTx (PR #439, merged). Happy to help with migration work here as well.

[zdcao121]

Hi @exopoiesis,
Thanks for the detailed context and for sharing your experience with CrystalFormer. At the moment, we don’t have plans to migrate from Haiku to Flax NNX. For our current use cases, dm-haiku==0.0.15 works well with jax==0.7.2.

Also, to clarify, dm-haiku is not archived but in maintenance mode, meaning development is primarily focused on bug fixes and ensuring compatibility with newer JAX releases.