Skip to main content

Ctrl+K

This site covers the old Flax Linen API. [Explore the new Flax NNX API ✨]

Quick start
Flax Basics
Guides
Examples
Glossary
Frequently Asked Questions (FAQ)
Developer notes
The Flax philosophy
How to contribute
API Reference
Flax NNX

.rst

Model inspection

Model inspection#

Model surgery
Extracting intermediate values
Extracting gradients of intermediate values

previous

Scale up Flax Modules on multiple devices

next

Model surgery

By The Flax authors

© Copyright 2023, The Flax authors.