Training Neural Networks with Manifold Constraints via JAX and FLAX#
In the following several examples, we mainly aim to illustrate that it is easy to describe neural networks with manifold constraints by the build-in layers in CDOpt based on JAX and FLAX. Therefore, training these networks can be greatly facilitated by the advanced features from JAX and FLAX and various optimizers provided by Optax packages.