Migrating# Convert PyTorch models to Flax FC Layers Convolutions Convolutions and FC Layers Batch Norm Average Pooling Transposed Convolutions NNX 0.10 to NNX 0.11 Using Rngs in NNX Transforms Loading Checkpoints with RNGs Optimizer Updates Pytrees containing NNX Objects Flax Linen to Flax NNX Basic Module definition Variable creation Training step and compilation Collections and variable types Using multiple methods Transformations Scan over layers Using TrainState in Flax NNX Haiku to Flax NNX Basic Module definition Variable creation Training step and compilation Handling non-parameter states Using multiple methods Transformations Scan over layers Top-level Haiku functions vs top-level Flax modules
