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.
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-haikuhas 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_shaperemoved fromShapedArray).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.