Deep learning with JAX and Flax

Book about JAX and Flax. Covering how to perform deep learning in JAX and Flax.

Topics include:

• Handling state in JAX & Flax,
• Transfer learning with JAX & Flax,
• Activation functions in JAX and Flax
• Optimizers in JAX and Flax
• Elegy(High-level API for deep learning in JAX & Flax)
• LSTM in JAX & Flax,
• JAX loss functions
• Distributed training with JAX & Flax
• How to load datasets in JAX with TensorFlow