Guides# Flax fundamentals JAX 101 Flax Basics Managing Parameters and State setup vs compact Dealing with Flax Module arguments Data preprocesing Processing the entire Dataset Training techniques Batch normalization Dropout Learning rate scheduling Transfer learning Save and load checkpoints Parallel training Ensembling on multiple devices Scale up Flax Modules on multiple devices with pjit Model inspection Model surgery Extracting intermediate values Extracting gradients of intermediate values Converting and upgrading Convert PyTorch models to Flax Migrate checkpointing to Orbax Upgrading my codebase to Optax Upgrading my codebase to Linen The Sharp Bits 🔪 flax.linen.Dropout layer and randomness
