Minor release to improve the ONNX exports.
Changed:
- All ONNX models are now exported using jax2onnx. The new exporter is better in many ways, so the resulting ONNX models look a little cleaner and are better supported by Netron. For now, the models themselves have stayed the same, the weights have been simply transplanted from the Keras models to the corresponding Flax NNX models.
My aim is to have all future models and experiments done in a pure JAX environment, without having to rely on Keras or TensorFlow. The remaining hurdle was the lack of a reliable ONNX exporter, but jax2onnx is mature enough now.