Haiku is a library for building neural networks in JAX.
conda install conda-forge::dm-haiku
Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet, a neural network library for TensorFlow.