CINXE.COM

<!DOCTYPE html> <html lang="en" data-content_root="" > <head> <meta charset="utf-8" /> <meta name="viewport" content="width=device-width, initial-scale=1.0" /> <title>Flax</title> <script data-cfasync="false"> document.documentElement.dataset.mode = localStorage.getItem("mode") || ""; document.documentElement.dataset.theme = localStorage.getItem("theme") || ""; </script>  <link href="_static/styles/theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="_static/styles/bootstrap.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="_static/styles/pydata-sphinx-theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="_static/vendor/fontawesome/6.5.2/css/all.min.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="_static/vendor/fontawesome/6.5.2/webfonts/fa-solid-900.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="_static/vendor/fontawesome/6.5.2/webfonts/fa-brands-400.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="_static/vendor/fontawesome/6.5.2/webfonts/fa-regular-400.woff2" /> <link rel="stylesheet" type="text/css" href="_static/pygments.css" /> <link rel="stylesheet" type="text/css" href="_static/styles/sphinx-book-theme.css" /> <link rel="stylesheet" type="text/css" href="_static/mystnb.4510f1fc1dee50b3e5859aac5469c37c29e427902b24a333a5f9fcb2f0b3ac41.css" /> <link rel="stylesheet" type="text/css" href="_static/sphinx-design.5ea377869091fd0449014c60fc090103.min.css" /> <link rel="stylesheet" type="text/css" href="_static/css/flax_theme.css" />  <link rel="preload" as="script" href="_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b" /> <link rel="preload" as="script" href="_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b" /> <script src="_static/vendor/fontawesome/6.5.2/js/all.min.js?digest=dfe6caa3a7d634c4db9b"></script> <script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script> <script src="_static/jquery.js"></script> <script src="_static/underscore.js"></script> <script src="_static/_sphinx_javascript_frameworks_compat.js"></script> <script src="_static/doctools.js"></script> <script src="_static/sphinx_highlight.js"></script> <script src="_static/scripts/sphinx-book-theme.js"></script> <script src="_static/design-tabs.js"></script> <script>DOCUMENTATION_OPTIONS.pagename = 'index';</script> <link rel="shortcut icon" href="_static/flax.png"/> <link rel="index" title="Index" href="genindex.html" /> <link rel="search" title="Search" href="search.html" /> <link rel="next" title="Flax basics" href="nnx_basics.html" /> <meta name="viewport" content="width=device-width, initial-scale=1"/> <meta name="docsearch:language" content="en"/> <script async type="text/javascript" src="/_/static/javascript/readthedocs-addons.js"></script><meta name="readthedocs-project-slug" content="flax" /><meta name="readthedocs-version-slug" content="latest" /><meta name="readthedocs-resolver-filename" content="/" /><meta name="readthedocs-http-status" content="200" /></head> <body data-bs-spy="scroll" data-bs-target=".bd-toc-nav" data-offset="180" data-bs-root-margin="0px 0px -60%" data-default-mode=""> <div id="pst-skip-link" class="skip-link d-print-none"><a href="#main-content">Skip to main content</a></div> <div id="pst-scroll-pixel-helper"></div> <button type="button" class="btn rounded-pill" id="pst-back-to-top"> <i class="fa-solid fa-arrow-up"></i>Back to top</button> <input type="checkbox" class="sidebar-toggle" id="pst-primary-sidebar-checkbox"/> <label class="overlay overlay-primary" for="pst-primary-sidebar-checkbox"></label> <input type="checkbox" class="sidebar-toggle" id="pst-secondary-sidebar-checkbox"/> <label class="overlay overlay-secondary" for="pst-secondary-sidebar-checkbox"></label> <div class="search-button__wrapper"> <div class="search-button__overlay"></div> <div class="search-button__search-container"> <form class="bd-search d-flex align-items-center" action="search.html" method="get"> <i class="fa-solid fa-magnifying-glass"></i> <input type="search" class="form-control" name="q" id="search-input" placeholder="Search..." aria-label="Search..." autocomplete="off" autocorrect="off" autocapitalize="off" spellcheck="false"/> <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd>K</kbd></span> </form></div> </div> <div class="pst-async-banner-revealer d-none"> <aside id="bd-header-version-warning" class="d-none d-print-none" aria-label="Version warning"></aside> </div> <aside class="bd-header-announcement" aria-label="Announcement"> <div class="bd-header-announcement__content"> <a href="https://flax-linen.readthedocs.io/en/latest" style="text-decoration: none; color: white;" > This site covers the new Flax NNX API. <span style="color: lightgray;">[Click here for the old <b>Flax Linen</b> API]</span> </a> </div> </aside> <header class="bd-header navbar navbar-expand-lg bd-navbar d-print-none"> </header> <div class="bd-container"> <div class="bd-container__inner bd-page-width"> <div class="bd-sidebar-primary bd-sidebar"> <div class="sidebar-header-items sidebar-primary__section"> </div> <div class="sidebar-primary-items__start sidebar-primary__section"> <div class="sidebar-primary-item"> <a class="navbar-brand logo" href="#"> <img src="_static/flax.png" class="logo__image only-light" alt=" - Home"/> <script>document.write(`<img src="_static/flax.png" class="logo__image only-dark" alt=" - Home"/>`);</script> </a></div> <div class="sidebar-primary-item"> <script> document.write(` <button class="btn search-button-field search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip"> <i class="fa-solid fa-magnifying-glass"></i> <span class="search-button__default-text">Search</span> <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd></span> </button> `); </script></div> <div class="sidebar-primary-item"><nav class="bd-links bd-docs-nav" aria-label="Main"> <div class="bd-toc-item navbar-nav active"> <ul class="nav bd-sidenav"> <li class="toctree-l1"><a class="reference internal" href="nnx_basics.html">Flax basics</a></li> <li class="toctree-l1"><a class="reference internal" href="mnist_tutorial.html">MNIST tutorial</a></li> <li class="toctree-l1"><a class="reference internal" href="why.html">Why Flax NNX?</a></li> <li class="toctree-l1 has-children"><a class="reference internal" href="guides/index.html">Guides</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l2"><a class="reference internal" href="guides/transforms.html">Transformations</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/flax_gspmd.html">Scale up on multiple devices</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/filters_guide.html">Using Filters</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/randomness.html">Randomness</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/performance.html">Performance considerations</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/linen_to_nnx.html">Evolution from Flax Linen to NNX</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/bridge_guide.html">Use Flax NNX and Linen together</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/surgery.html">Model surgery</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/checkpointing.html">Save and load checkpoints</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/jax_and_nnx_transforms.html">Flax NNX vs JAX transformations</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/haiku_to_flax.html">Migrating from Haiku to Flax</a></li> <li class="toctree-l2"><a class="reference internal" href="guides/gemma.html">Example: Using pretrained Gemma for inference with Flax NNX</a></li> </ul> </details></li> <li class="toctree-l1 has-children"><a class="reference internal" href="examples/index.html">Examples</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l2"><a class="reference internal" href="examples/core_examples.html">Core examples</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference internal" href="nnx_glossary.html">Flax NNX glossary</a></li> <li class="toctree-l1"><a class="reference external" href="https://flax.readthedocs.io/en/latest/philosophy.html">The Flax philosophy</a></li> <li class="toctree-l1"><a class="reference external" href="https://flax.readthedocs.io/en/latest/contributing.html">How to contribute</a></li> <li class="toctree-l1 has-children"><a class="reference internal" href="api_reference/index.html">API Reference</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l2 has-children"><a class="reference internal" href="api_reference/flax.nnx/index.html">flax.nnx</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/graph.html">graph</a></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/module.html">module</a></li> <li class="toctree-l3 has-children"><a class="reference internal" href="api_reference/flax.nnx/nn/index.html">nn</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/nn/activations.html">Activation functions</a></li> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/nn/attention.html">Attention</a></li> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/nn/dtypes.html">Dtypes</a></li> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/nn/initializers.html">Initializers</a></li> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/nn/linear.html">Linear</a></li> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/nn/lora.html">LoRA</a></li> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/nn/normalization.html">Normalization</a></li> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/nn/recurrent.html">Recurrent</a></li> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/nn/stochastic.html">Stochastic</a></li> </ul> </details></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/rnglib.html">rnglib</a></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/spmd.html">spmd</a></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/state.html">state</a></li> <li class="toctree-l3 has-children"><a class="reference internal" href="api_reference/flax.nnx/training/index.html">training</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/training/metrics.html">Metrics</a></li> <li class="toctree-l4"><a class="reference internal" href="api_reference/flax.nnx/training/optimizer.html">Optimizer</a></li> </ul> </details></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/transforms.html">transforms</a></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/variables.html">variables</a></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/helpers.html">helpers</a></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/visualization.html">visualization</a></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/filterlib.html">filterlib</a></li> <li class="toctree-l3"><a class="reference internal" href="api_reference/flax.nnx/bridge.html">bridge</a></li> </ul> </details></li> <li class="toctree-l2"><a class="reference internal" href="api_reference/flax.core.frozen_dict.html">flax.core.frozen_dict package</a></li> <li class="toctree-l2"><a class="reference internal" href="api_reference/flax.struct.html">flax.struct package</a></li> <li class="toctree-l2"><a class="reference internal" href="api_reference/flax.training.html">flax.training package</a></li> </ul> </details></li> </ul> </div> </nav></div> </div> <div class="sidebar-primary-items__end sidebar-primary__section"> </div> <div id="rtd-footer-container"></div> </div> <main id="main-content" class="bd-main" role="main"> <div class="sbt-scroll-pixel-helper"></div> <div class="bd-content"> <div class="bd-article-container"> <div class="bd-header-article d-print-none"> <div class="header-article-items header-article__inner"> <div class="header-article-items__start"> <div class="header-article-item"><button class="sidebar-toggle primary-toggle btn btn-sm" title="Toggle primary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip"> <span class="fa-solid fa-bars"></span> </button></div> </div> <div class="header-article-items__end"> <div class="header-article-item"> <div class="article-header-buttons"> <a href="https://github.com/google/flax" target="_blank" class="btn btn-sm btn-source-repository-button" title="Source repository" data-bs-placement="bottom" data-bs-toggle="tooltip" > <span class="btn__icon-container"> <i class="fab fa-github"></i> </span> </a> <div class="dropdown dropdown-download-buttons"> <button class="btn dropdown-toggle" type="button" data-bs-toggle="dropdown" aria-expanded="false" aria-label="Download this page"> <i class="fas fa-download"></i> </button> <ul class="dropdown-menu"> <li><a href="_sources/index.rst" target="_blank" class="btn btn-sm btn-download-source-button dropdown-item" title="Download source file" data-bs-placement="left" data-bs-toggle="tooltip" > <span class="btn__icon-container"> <i class="fas fa-file"></i> </span> <span class="btn__text-container">.rst</span> </a> </li> <li> <button onclick="window.print()" class="btn btn-sm btn-download-pdf-button dropdown-item" title="Print to PDF" data-bs-placement="left" data-bs-toggle="tooltip" > <span class="btn__icon-container"> <i class="fas fa-file-pdf"></i> </span> <span class="btn__text-container">.pdf</span> </button> </li> </ul> </div> <button onclick="toggleFullScreen()" class="btn btn-sm btn-fullscreen-button" title="Fullscreen mode" data-bs-placement="bottom" data-bs-toggle="tooltip" > <span class="btn__icon-container"> <i class="fas fa-expand"></i> </span> </button> <script> document.write(` <button class="btn btn-sm nav-link pst-navbar-icon theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip"> <i class="theme-switch fa-solid fa-sun fa-lg" data-mode="light"></i> <i class="theme-switch fa-solid fa-moon fa-lg" data-mode="dark"></i> <i class="theme-switch fa-solid fa-circle-half-stroke fa-lg" data-mode="auto"></i> </button> `); </script> <script> document.write(` <button class="btn btn-sm pst-navbar-icon search-button search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip"> <i class="fa-solid fa-magnifying-glass fa-lg"></i> </button> `); </script> <button class="sidebar-toggle secondary-toggle btn btn-sm" title="Toggle secondary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip"> <span class="fa-solid fa-list"></span> </button> </div></div> </div> </div> </div> <div id="jb-print-docs-body" class="onlyprint"> <h1>Flax</h1>  <div id="print-main-content"> <div id="jb-print-toc"> <div> <h2> Contents </h2> </div> <nav aria-label="Page"> <ul class="visible nav section-nav flex-column"> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#features">Features</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#basic-usage">Basic usage</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#installation">Installation</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#learn-more">Learn more</a></li> </ul> </nav> </div> </div> </div> <div id="searchbox"></div> <article class="bd-article"> <div class="section" id="flax"> <h1>Flax<a class="headerlink" href="#flax" title="Permalink to this heading">#</a></h1> <div class="sd-text-left sd-font-italic docutils"> <p><strong>N</strong>eural <strong>N</strong>etworks for JA<strong>X</strong></p> </div> <hr class="docutils" /> <p>Flax provides a <strong>flexible end-to-end user experience for researchers and developers who use JAX for neural networks</strong>. Flax enables you to use the full power of <a class="reference external" href="https://jax.readthedocs.io">JAX</a>.</p> <p>At the core of Flax is <strong>NNX - a simplified API that makes it easier to create, inspect, debug, and analyze neural networks in JAX.</strong> Flax NNX has first class support for Python reference semantics, enabling users to express their models using regular Python objects. Flax NNX is an evolution of the previous <a class="reference external" href="https://flax-linen.readthedocs.io/">Flax Linen</a> API, and it took years of experience to bring a simpler and more user-friendly API.</p> <div class="admonition note"> <p class="admonition-title">Note</p> <p>Flax Linen API is not going to be deprecated in the near future as most of Flax users still rely on this API. However, new users are encouraged to use Flax NNX. Check out <a class="reference external" href="why.html">Why Flax NNX</a> for a comparison between Flax NNX and Linen, and our reasoning to make the new API.</p> <p>To move your Flax Linen codebase to Flax NNX, get familiarized with the API in <a class="reference external" href="https://flax.readthedocs.io/en/latest/nnx_basics.html">NNX Basics</a> and then start your move following the <a class="reference external" href="guides/linen_to_nnx.html">evolution guide</a>.</p> </div> <div class="section" id="features"> <h2>Features<a class="headerlink" href="#features" title="Permalink to this heading">#</a></h2> <div class="sd-container-fluid sd-sphinx-override sd-mb-4 docutils"> <div class="sd-row docutils"> <div class="sd-col sd-d-flex-column sd-col-12 sd-col-xs-12 sd-col-sm-12 sd-col-md-12 sd-col-lg-6 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-none sd-border-0 docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold sd-fs-5 docutils"> Pythonic</div> <div class="sd-font-normal docutils"> <p class="sd-card-text">Flax NNX supports the use of regular Python objects, providing an intuitive and predictable development experience.</p> </div> </div> </div> </div> <div class="sd-col sd-d-flex-column sd-col-12 sd-col-xs-12 sd-col-sm-12 sd-col-md-12 sd-col-lg-6 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-none sd-border-0 docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold sd-fs-5 docutils"> Simple</div> <div class="sd-font-normal docutils"> <p class="sd-card-text">Flax NNX relies on Python鈥檚 object model, which results in simplicity for the user and increases development speed.</p> </div> </div> </div> </div> <div class="sd-col sd-d-flex-column sd-col-12 sd-col-xs-12 sd-col-sm-12 sd-col-md-12 sd-col-lg-6 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-none sd-border-0 docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold sd-fs-5 docutils"> Expressive</div> <div class="sd-font-normal docutils"> <p class="sd-card-text">Flax NNX allows fine-grained control of the model鈥檚 state via its <a class="reference external" href="https://flax.readthedocs.io/en/latest/nnx/filters_guide.html">Filter</a> system.</p> </div> </div> </div> </div> <div class="sd-col sd-d-flex-column sd-col-12 sd-col-xs-12 sd-col-sm-12 sd-col-md-12 sd-col-lg-6 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-none sd-border-0 docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold sd-fs-5 docutils"> Familiar</div> <div class="sd-font-normal docutils"> <p class="sd-card-text">Flax NNX makes it very easy to integrate objects with regular JAX code via the <a class="reference external" href="nnx_basics.html#the-flax-functional-api">Functional API</a>.</p> </div> </div> </div> </div> </div> </div> </div> <div class="section" id="basic-usage"> <h2>Basic usage<a class="headerlink" href="#basic-usage" title="Permalink to this heading">#</a></h2> <div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">flax</span> <span class="kn">import</span> <span class="n">nnx</span> <span class="kn">import</span> <span class="nn">optax</span> <span class="k">class</span> <span class="nc">Model</span><span class="p">(</span><span class="n">nnx</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span> <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">din</span><span class="p">,</span> <span class="n">dmid</span><span class="p">,</span> <span class="n">dout</span><span class="p">,</span> <span class="n">rngs</span><span class="p">:</span> <span class="n">nnx</span><span class="o">.</span><span class="n">Rngs</span><span class="p">):</span> <span class="bp">self</span><span class="o">.</span><span class="n">linear</span> <span class="o">=</span> <span class="n">nnx</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">din</span><span class="p">,</span> <span class="n">dmid</span><span class="p">,</span> <span class="n">rngs</span><span class="o">=</span><span class="n">rngs</span><span class="p">)</span> <span class="bp">self</span><span class="o">.</span><span class="n">bn</span> <span class="o">=</span> <span class="n">nnx</span><span class="o">.</span><span class="n">BatchNorm</span><span class="p">(</span><span class="n">dmid</span><span class="p">,</span> <span class="n">rngs</span><span class="o">=</span><span class="n">rngs</span><span class="p">)</span> <span class="bp">self</span><span class="o">.</span><span class="n">dropout</span> <span class="o">=</span> <span class="n">nnx</span><span class="o">.</span><span class="n">Dropout</span><span class="p">(</span><span class="mf">0.2</span><span class="p">,</span> <span class="n">rngs</span><span class="o">=</span><span class="n">rngs</span><span class="p">)</span> <span class="bp">self</span><span class="o">.</span><span class="n">linear_out</span> <span class="o">=</span> <span class="n">nnx</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">dmid</span><span class="p">,</span> <span class="n">dout</span><span class="p">,</span> <span class="n">rngs</span><span class="o">=</span><span class="n">rngs</span><span class="p">)</span> <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span> <span class="n">x</span> <span class="o">=</span> <span class="n">nnx</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dropout</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">bn</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">linear</span><span class="p">(</span><span class="n">x</span><span class="p">))))</span> <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">linear_out</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="n">model</span> <span class="o">=</span> <span class="n">Model</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">64</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">rngs</span><span class="o">=</span><span class="n">nnx</span><span class="o">.</span><span class="n">Rngs</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span> <span class="c1"># eager initialization</span> <span class="n">optimizer</span> <span class="o">=</span> <span class="n">nnx</span><span class="o">.</span><span class="n">Optimizer</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">optax</span><span class="o">.</span><span class="n">adam</span><span class="p">(</span><span class="mf">1e-3</span><span class="p">))</span> <span class="c1"># reference sharing</span> <span class="nd">@nnx</span><span class="o">.</span><span class="n">jit</span> <span class="c1"># automatic state management for JAX transforms</span> <span class="k">def</span> <span class="nf">train_step</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">optimizer</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span> <span class="k">def</span> <span class="nf">loss_fn</span><span class="p">(</span><span class="n">model</span><span class="p">):</span> <span class="n">y_pred</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="c1"># call methods directly</span> <span class="k">return</span> <span class="p">((</span><span class="n">y_pred</span> <span class="o">-</span> <span class="n">y</span><span class="p">)</span> <span class="o">**</span> <span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span> <span class="n">loss</span><span class="p">,</span> <span class="n">grads</span> <span class="o">=</span> <span class="n">nnx</span><span class="o">.</span><span class="n">value_and_grad</span><span class="p">(</span><span class="n">loss_fn</span><span class="p">)(</span><span class="n">model</span><span class="p">)</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">grads</span><span class="p">)</span> <span class="c1"># in-place updates</span> <span class="k">return</span> <span class="n">loss</span> </pre></div> </div> </div> <div class="section" id="installation"> <h2>Installation<a class="headerlink" href="#installation" title="Permalink to this heading">#</a></h2> <p>Install via pip:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>pip<span class="w"> </span>install<span class="w"> </span>flax </pre></div> </div> <p>Or install the latest version from the repository:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>pip<span class="w"> </span>install<span class="w"> </span>git+https://github.com/google/flax.git </pre></div> </div> </div> <hr class="docutils" /> <div class="section" id="learn-more"> <h2>Learn more<a class="headerlink" href="#learn-more" title="Permalink to this heading">#</a></h2> <div class="sd-container-fluid sd-sphinx-override sd-mb-4 docutils"> <div class="sd-row docutils"> <div class="sd-col sd-d-flex-column sd-col-6 sd-col-xs-6 sd-col-sm-6 sd-col-md-6 sd-col-lg-4 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-sm sd-card-hover sd-text-black sd-bg-light docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold docutils"> <svg version="4.0.0.63c5cb3" width="2.0em" height="2.0em" class="sd-material-icon sd-material-icon-rocket_launch" viewBox="0 0 24 24" aria-hidden="true"><g><rect fill="none" height="24" width="24"></rect></g><g><g><path d="M9.19,6.35c-2.04,2.29-3.44,5.58-3.57,5.89L2,10.69l4.05-4.05c0.47-0.47,1.15-0.68,1.81-0.55L9.19,6.35L9.19,6.35z M11.17,17c0,0,3.74-1.55,5.89-3.7c5.4-5.4,4.5-9.62,4.21-10.57c-0.95-0.3-5.17-1.19-10.57,4.21C8.55,9.09,7,12.83,7,12.83 L11.17,17z M17.65,14.81c-2.29,2.04-5.58,3.44-5.89,3.57L13.31,22l4.05-4.05c0.47-0.47,0.68-1.15,0.55-1.81L17.65,14.81 L17.65,14.81z M9,18c0,0.83-0.34,1.58-0.88,2.12C6.94,21.3,2,22,2,22s0.7-4.94,1.88-6.12C4.42,15.34,5.17,15,6,15 C7.66,15,9,16.34,9,18z M13,9c0-1.1,0.9-2,2-2s2,0.9,2,2s-0.9,2-2,2S13,10.1,13,9z"></path></g></g></svg> Flax NNX Basics</div> </div> <a class="sd-stretched-link sd-hide-link-text reference external" href="nnx_basics.html"><span>nnx_basics.html</span></a></div> </div> <div class="sd-col sd-d-flex-column sd-col-6 sd-col-xs-6 sd-col-sm-6 sd-col-md-6 sd-col-lg-4 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-sm sd-card-hover sd-text-black sd-bg-light docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold docutils"> <svg version="4.0.0.63c5cb3" width="2.0em" height="2.0em" class="sd-material-icon sd-material-icon-library_books" viewBox="0 0 24 24" aria-hidden="true"><path d="M0 0h24v24H0z" fill="none"></path><path d="M4 6H2v14c0 1.1.9 2 2 2h14v-2H4V6zm16-4H8c-1.1 0-2 .9-2 2v12c0 1.1.9 2 2 2h12c1.1 0 2-.9 2-2V4c0-1.1-.9-2-2-2zm-1 9H9V9h10v2zm-4 4H9v-2h6v2zm4-8H9V5h10v2z"></path></svg> MNIST Tutorial</div> </div> <a class="sd-stretched-link sd-hide-link-text reference external" href="mnist_tutorial.html"><span>mnist_tutorial.html</span></a></div> </div> <div class="sd-col sd-d-flex-column sd-col-6 sd-col-xs-6 sd-col-sm-6 sd-col-md-6 sd-col-lg-4 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-sm sd-card-hover sd-text-black sd-bg-light docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold docutils"> <svg version="4.0.0.63c5cb3" width="2.0em" height="2.0em" class="sd-material-icon sd-material-icon-library_books" viewBox="0 0 24 24" aria-hidden="true"><path d="M0 0h24v24H0z" fill="none"></path><path d="M4 6H2v14c0 1.1.9 2 2 2h14v-2H4V6zm16-4H8c-1.1 0-2 .9-2 2v12c0 1.1.9 2 2 2h12c1.1 0 2-.9 2-2V4c0-1.1-.9-2-2-2zm-1 9H9V9h10v2zm-4 4H9v-2h6v2zm4-8H9V5h10v2z"></path></svg> Guides</div> </div> <a class="sd-stretched-link sd-hide-link-text reference external" href="guides/index.html"><span>guides/index.html</span></a></div> </div> <div class="sd-col sd-d-flex-column sd-col-6 sd-col-xs-6 sd-col-sm-6 sd-col-md-6 sd-col-lg-4 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-sm sd-card-hover sd-text-black sd-bg-light docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold docutils"> <svg version="4.0.0.63c5cb3" width="2.0em" height="2.0em" class="sd-material-icon sd-material-icon-transform" viewBox="0 0 24 24" aria-hidden="true"><path d="M0 0h24v24H0z" fill="none"></path><path d="M22 18v-2H8V4h2L7 1 4 4h2v2H2v2h4v8c0 1.1.9 2 2 2h8v2h-2l3 3 3-3h-2v-2h4zM10 8h6v6h2V8c0-1.1-.9-2-2-2h-6v2z"></path></svg> Flax Linen to Flax NNX</div> </div> <a class="sd-stretched-link sd-hide-link-text reference external" href="guides/linen_to_nnx.html"><span>guides/linen_to_nnx.html</span></a></div> </div> <div class="sd-col sd-d-flex-column sd-col-6 sd-col-xs-6 sd-col-sm-6 sd-col-md-6 sd-col-lg-4 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-sm sd-card-hover sd-text-black sd-bg-light docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold docutils"> <svg version="4.0.0.63c5cb3" width="2.0em" height="2.0em" class="sd-material-icon sd-material-icon-menu_book" viewBox="0 0 24 24" aria-hidden="true"><g><rect fill="none" height="24" width="24"></rect></g><g><g></rect><g><path d="M21,5c-1.11-0.35-2.33-0.5-3.5-0.5c-1.95,0-4.05,0.4-5.5,1.5c-1.45-1.1-3.55-1.5-5.5-1.5S2.45,4.9,1,6v14.65 c0,0.25,0.25,0.5,0.5,0.5c0.1,0,0.15-0.05,0.25-0.05C3.1,20.45,5.05,20,6.5,20c1.95,0,4.05,0.4,5.5,1.5c1.35-0.85,3.8-1.5,5.5-1.5 c1.65,0,3.35,0.3,4.75,1.05c0.1,0.05,0.15,0.05,0.25,0.05c0.25,0,0.5-0.25,0.5-0.5V6C22.4,5.55,21.75,5.25,21,5z M21,18.5 c-1.1-0.35-2.3-0.5-3.5-0.5c-1.7,0-4.15,0.65-5.5,1.5V8c1.35-0.85,3.8-1.5,5.5-1.5c1.2,0,2.4,0.15,3.5,0.5V18.5z"></path><g><path d="M17.5,10.5c0.88,0,1.73,0.09,2.5,0.26V9.24C19.21,9.09,18.36,9,17.5,9c-1.7,0-3.24,0.29-4.5,0.83v1.66 C14.13,10.85,15.7,10.5,17.5,10.5z"></path><path d="M13,12.49v1.66c1.13-0.64,2.7-0.99,4.5-0.99c0.88,0,1.73,0.09,2.5,0.26V11.9c-0.79-0.15-1.64-0.24-2.5-0.24 C15.8,11.66,14.26,11.96,13,12.49z"></path><path d="M17.5,14.33c-1.7,0-3.24,0.29-4.5,0.83v1.66c1.13-0.64,2.7-0.99,4.5-0.99c0.88,0,1.73,0.09,2.5,0.26v-1.52 C19.21,14.41,18.36,14.33,17.5,14.33z"></path></g></g></g></svg> API reference</div> </div> <a class="sd-stretched-link sd-hide-link-text reference external" href="api_reference/index.html"><span>api_reference/index.html</span></a></div> </div> <div class="sd-col sd-d-flex-column sd-col-6 sd-col-xs-6 sd-col-sm-6 sd-col-md-6 sd-col-lg-4 docutils"> <div class="sd-card sd-sphinx-override sd-mb-3 sd-shadow-sm sd-card-hover sd-text-black sd-bg-light docutils"> <div class="sd-card-body docutils"> <div class="sd-card-title sd-font-weight-bold docutils"> <svg version="4.0.0.63c5cb3" width="2.0em" height="2.0em" class="sd-material-icon sd-material-icon-import_contacts" viewBox="0 0 24 24" aria-hidden="true"><g><rect fill="none" height="24" width="24"></rect></g><g><g><g><path d="M17.5,4.5c-1.95,0-4.05,0.4-5.5,1.5c-1.45-1.1-3.55-1.5-5.5-1.5S2.45,4.9,1,6v14.65c0,0.65,0.73,0.45,0.75,0.45 C3.1,20.45,5.05,20,6.5,20c1.95,0,4.05,0.4,5.5,1.5c1.35-0.85,3.8-1.5,5.5-1.5c1.65,0,3.35,0.3,4.75,1.05 C22.66,21.26,23,20.86,23,20.6V6C21.51,4.88,19.37,4.5,17.5,4.5z M21,18.5c-1.1-0.35-2.3-0.5-3.5-0.5c-1.7,0-4.15,0.65-5.5,1.5V8 c1.35-0.85,3.8-1.5,5.5-1.5c1.2,0,2.4,0.15,3.5,0.5V18.5z"></path></g></g></g></svg> Glossary</div> </div> <a class="sd-stretched-link sd-hide-link-text reference external" href="nnx_glossary.html"><span>nnx_glossary.html</span></a></div> </div> </div> </div> <hr class="docutils" /> <div class="toctree-wrapper compound"> </div> </div> </div> </article> <footer class="prev-next-footer d-print-none"> <div class="prev-next-area"> <a class="right-next" href="nnx_basics.html" title="next page"> <div class="prev-next-info"> <p class="prev-next-subtitle">next</p> <p class="prev-next-title">Flax basics</p> </div> <i class="fa-solid fa-angle-right"></i> </a> </div> </footer> </div> <div class="bd-sidebar-secondary bd-toc"><div class="sidebar-secondary-items sidebar-secondary__inner"> <div class="sidebar-secondary-item"> <div class="page-toc tocsection onthispage"> <i class="fa-solid fa-list"></i> Contents </div> <nav class="bd-toc-nav page-toc"> <ul class="visible nav section-nav flex-column"> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#features">Features</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#basic-usage">Basic usage</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#installation">Installation</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#learn-more">Learn more</a></li> </ul> </nav></div> </div></div> </div> <footer class="bd-footer-content"> <div class="bd-footer-content__inner container"> <div class="footer-item"> <p class="component-author"> By The Flax authors </p> </div> <div class="footer-item"> <p class="copyright"> 漏 Copyright 2023, The Flax authors. <br/> </p> </div> <div class="footer-item"> </div> <div class="footer-item"> </div> </div> </footer> </main> </div> </div>  <script src="_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b"></script> <script src="_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b"></script> <footer class="bd-footer"> </footer> </body> </html>