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.linen</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 = 'api_reference/flax.linen/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="Module" href="module.html" /> <link rel="prev" title="flax.jax_utils package" href="../flax.jax_utils.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-linen" /><meta name="readthedocs-version-slug" content="latest" /><meta name="readthedocs-resolver-filename" content="/api_reference/flax.linen/index.html" /><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.readthedocs.io/en/latest/index.html" style="text-decoration: none; color: white;" > This site covers the old Flax Linen API. <span style="color: lightgray;">[Explore the new <b>Flax NNX</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="../../index.html"> <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="current nav bd-sidenav"> <li class="toctree-l1"><a class="reference internal" href="../../quick_start.html">Quick start</a></li> <li class="toctree-l1"><a class="reference internal" href="../../guides/flax_fundamentals/flax_basics.html">Flax Basics</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 has-children"><a class="reference internal" href="../../guides/flax_fundamentals/index.html">Flax fundamentals</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 external" href="https://jax.readthedocs.io/en/latest/jax-101/index.html">JAX 101</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/flax_fundamentals/flax_basics.html">Flax Basics</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/flax_fundamentals/state_params.html">Managing Parameters and State</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/flax_fundamentals/setup_or_nncompact.html"><code class="docutils literal notranslate"><span class="pre">setup</span></code> vs <code class="docutils literal notranslate"><span class="pre">compact</span></code></a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/flax_fundamentals/arguments.html">Dealing with Flax Module arguments</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/flax_fundamentals/rng_guide.html">Randomness and PRNGs in Flax</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../../guides/data_preprocessing/index.html">Data preprocessing</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="../../guides/data_preprocessing/full_eval.html">Processing the entire Dataset</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/data_preprocessing/loading_datasets.html">Loading datasets</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../../guides/training_techniques/index.html">Training techniques</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="../../guides/training_techniques/batch_norm.html">Batch normalization</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/training_techniques/dropout.html">Dropout</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/training_techniques/lr_schedule.html">Learning rate scheduling</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/training_techniques/transfer_learning.html">Transfer learning</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/training_techniques/use_checkpointing.html">Save and load checkpoints</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../../guides/parallel_training/index.html">Parallel training</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="../../guides/parallel_training/ensembling.html">Ensembling on multiple devices</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/parallel_training/flax_on_pjit.html">Scale up Flax Modules on multiple devices</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../../guides/model_inspection/index.html">Model inspection</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="../../guides/model_inspection/model_surgery.html">Model surgery</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/model_inspection/extracting_intermediates.html">Extracting intermediate values</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../../guides/converting_and_upgrading/index.html">Converting and upgrading</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="../../guides/converting_and_upgrading/haiku_migration_guide.html">Migrating from Haiku to Flax</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/converting_and_upgrading/convert_pytorch_to_flax.html">Convert PyTorch models to Flax</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/converting_and_upgrading/orbax_upgrade_guide.html">Migrate checkpointing to Orbax</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/converting_and_upgrading/optax_update_guide.html">Upgrading my codebase to Optax</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/converting_and_upgrading/linen_upgrade_guide.html">Upgrading my codebase to Linen</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/converting_and_upgrading/rnncell_upgrade_guide.html">RNNCellBase Upgrade Guide</a></li> <li class="toctree-l3"><a class="reference internal" href="../../guides/converting_and_upgrading/regular_dict_upgrade_guide.html">Migrate to regular dicts</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../../guides/quantization/index.html">Quantization</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="../../guides/quantization/fp8_basics.html">User Guide on Using FP8</a></li> </ul> </details></li> <li class="toctree-l2"><a class="reference internal" href="../../guides/flax_sharp_bits.html">The Sharp Bits</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> <li class="toctree-l2"><a class="reference internal" href="../../examples/google_research_examples.html">Google Research examples</a></li> <li class="toctree-l2"><a class="reference internal" href="../../examples/repositories_that_use_flax.html">Repositories that use Flax</a></li> <li class="toctree-l2"><a class="reference internal" href="../../examples/community_examples.html">Community examples</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference internal" href="../../glossary.html">Glossary</a></li> <li class="toctree-l1"><a class="reference internal" href="../../faq.html">Frequently Asked Questions (FAQ)</a></li> <li class="toctree-l1 has-children"><a class="reference internal" href="../../developer_notes/index.html">Developer notes</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="../../developer_notes/module_lifecycle.html">The Flax Module lifecycle</a></li> <li class="toctree-l2"><a class="reference internal" href="../../developer_notes/lift.html">Lifted transformations</a></li> <li class="toctree-l2"><a class="reference external" href="https://github.com/google/flax/tree/main/docs/flip">FLIPs</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference internal" href="../../philosophy.html">The Flax philosophy</a></li> <li class="toctree-l1"><a class="reference internal" href="../../contributing.html">How to contribute</a></li> <li class="toctree-l1 current active has-children"><a class="reference internal" href="../index.html">API Reference</a><details open="open"><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul class="current"> <li class="toctree-l2"><a class="reference internal" href="../flax.config.html">flax.config package</a></li> <li class="toctree-l2"><a class="reference internal" href="../flax.core.frozen_dict.html">flax.core.frozen_dict package</a></li> <li class="toctree-l2"><a class="reference internal" href="../flax.cursor.html">flax.cursor package</a></li> <li class="toctree-l2"><a class="reference internal" href="../flax.errors.html">flax.errors package</a></li> <li class="toctree-l2"><a class="reference internal" href="../flax.jax_utils.html">flax.jax_utils package</a></li> <li class="toctree-l2 current active has-children"><a class="current reference internal" href="#">flax.linen</a><details open="open"><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="module.html">Module</a></li> <li class="toctree-l3"><a class="reference internal" href="init_apply.html">Init/Apply</a></li> <li class="toctree-l3"><a class="reference internal" href="layers.html">Layers</a></li> <li class="toctree-l3"><a class="reference internal" href="activation_functions.html">Activation functions</a></li> <li class="toctree-l3"><a class="reference internal" href="initializers.html">Initializers</a></li> <li class="toctree-l3"><a class="reference internal" href="transformations.html">Transformations</a></li> <li class="toctree-l3"><a class="reference internal" href="inspection.html">Inspection</a></li> <li class="toctree-l3"><a class="reference internal" href="variable.html">Variable dictionary</a></li> <li class="toctree-l3"><a class="reference internal" href="spmd.html">SPMD</a></li> <li class="toctree-l3"><a class="reference internal" href="decorators.html">Decorators</a></li> <li class="toctree-l3"><a class="reference internal" href="profiling.html">Profiling</a></li> </ul> </details></li> <li class="toctree-l2"><a class="reference internal" href="../flax.serialization.html">flax.serialization package</a></li> <li class="toctree-l2"><a class="reference internal" href="../flax.struct.html">flax.struct package</a></li> <li class="toctree-l2"><a class="reference internal" href="../flax.traceback_util.html">flax.traceback_util package</a></li> <li class="toctree-l2"><a class="reference internal" href="../flax.training.html">flax.training package</a></li> <li class="toctree-l2"><a class="reference internal" href="../flax.traverse_util.html">flax.traverse_util package</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference external" href="https://flax.readthedocs.io/en/latest/index.html">Flax NNX</a></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/api_reference/flax.linen/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> </div></div> </div> </div> </div> <div id="jb-print-docs-body" class="onlyprint"> <h1>flax.linen</h1>  <div id="print-main-content"> <div id="jb-print-toc"> </div> </div> </div> <div id="searchbox"></div> <article class="bd-article"> <div class="section" id="flax-linen"> <h1>flax.linen<a class="headerlink" href="#flax-linen" title="Permalink to this heading">#</a></h1> <p>Linen is the Flax Module system. Read more about our design goals in the <a class="reference external" href="https://github.com/google/flax/blob/main/flax/linen/README.md">Linen README</a>.</p> <div class="toctree-wrapper compound"> <ul> <li class="toctree-l1"><a class="reference internal" href="module.html">Module</a><ul> <li class="toctree-l2"><a class="reference internal" href="module.html#flax.linen.Module"><code class="docutils literal notranslate"><span class="pre">Module</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="module.html#flax.linen.apply"><code class="docutils literal notranslate"><span class="pre">apply()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="module.html#flax.linen.init"><code class="docutils literal notranslate"><span class="pre">init()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="module.html#flax.linen.init_with_output"><code class="docutils literal notranslate"><span class="pre">init_with_output()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="module.html#flax.linen.intercept_methods"><code class="docutils literal notranslate"><span class="pre">intercept_methods()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="module.html#flax.linen.share_scope"><code class="docutils literal notranslate"><span class="pre">share_scope()</span></code></a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="init_apply.html">Init/Apply</a><ul> <li class="toctree-l2"><a class="reference internal" href="init_apply.html#flax.linen.apply"><code class="docutils literal notranslate"><span class="pre">apply()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="init_apply.html#flax.linen.init"><code class="docutils literal notranslate"><span class="pre">init()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="init_apply.html#flax.linen.init_with_output"><code class="docutils literal notranslate"><span class="pre">init_with_output()</span></code></a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="layers.html">Layers</a><ul> <li class="toctree-l2"><a class="reference internal" href="layers.html#linear-modules">Linear Modules</a></li> <li class="toctree-l2"><a class="reference internal" href="layers.html#pooling">Pooling</a></li> <li class="toctree-l2"><a class="reference internal" href="layers.html#normalization">Normalization</a></li> <li class="toctree-l2"><a class="reference internal" href="layers.html#combinators">Combinators</a></li> <li class="toctree-l2"><a class="reference internal" href="layers.html#stochastic">Stochastic</a></li> <li class="toctree-l2"><a class="reference internal" href="layers.html#attention">Attention</a></li> <li class="toctree-l2"><a class="reference internal" href="layers.html#recurrent">Recurrent</a></li> <li class="toctree-l2"><a class="reference internal" href="layers.html#batchapply">BatchApply</a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="activation_functions.html">Activation functions</a><ul> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.PReLU"><code class="docutils literal notranslate"><span class="pre">PReLU</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.celu"><code class="docutils literal notranslate"><span class="pre">celu()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.elu"><code class="docutils literal notranslate"><span class="pre">elu()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.gelu"><code class="docutils literal notranslate"><span class="pre">gelu()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.glu"><code class="docutils literal notranslate"><span class="pre">glu()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.hard_sigmoid"><code class="docutils literal notranslate"><span class="pre">hard_sigmoid()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.hard_silu"><code class="docutils literal notranslate"><span class="pre">hard_silu()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.hard_swish"><code class="docutils literal notranslate"><span class="pre">hard_swish()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.hard_tanh"><code class="docutils literal notranslate"><span class="pre">hard_tanh()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.leaky_relu"><code class="docutils literal notranslate"><span class="pre">leaky_relu()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.log_sigmoid"><code class="docutils literal notranslate"><span class="pre">log_sigmoid()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.log_softmax"><code class="docutils literal notranslate"><span class="pre">log_softmax()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.logsumexp"><code class="docutils literal notranslate"><span class="pre">logsumexp()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.one_hot"><code class="docutils literal notranslate"><span class="pre">one_hot()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.relu"><code class="docutils literal notranslate"><span class="pre">relu()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.selu"><code class="docutils literal notranslate"><span class="pre">selu()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.sigmoid"><code class="docutils literal notranslate"><span class="pre">sigmoid()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.silu"><code class="docutils literal notranslate"><span class="pre">silu()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.soft_sign"><code class="docutils literal notranslate"><span class="pre">soft_sign()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.softmax"><code class="docutils literal notranslate"><span class="pre">softmax()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.softplus"><code class="docutils literal notranslate"><span class="pre">softplus()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.standardize"><code class="docutils literal notranslate"><span class="pre">standardize()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.swish"><code class="docutils literal notranslate"><span class="pre">swish()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="activation_functions.html#flax.linen.activation.tanh"><code class="docutils literal notranslate"><span class="pre">tanh()</span></code></a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="initializers.html">Initializers</a><ul> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.constant"><code class="docutils literal notranslate"><span class="pre">constant()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.delta_orthogonal"><code class="docutils literal notranslate"><span class="pre">delta_orthogonal()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.glorot_normal"><code class="docutils literal notranslate"><span class="pre">glorot_normal()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.glorot_uniform"><code class="docutils literal notranslate"><span class="pre">glorot_uniform()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.he_normal"><code class="docutils literal notranslate"><span class="pre">he_normal()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.he_uniform"><code class="docutils literal notranslate"><span class="pre">he_uniform()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.kaiming_normal"><code class="docutils literal notranslate"><span class="pre">kaiming_normal()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.kaiming_uniform"><code class="docutils literal notranslate"><span class="pre">kaiming_uniform()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.lecun_normal"><code class="docutils literal notranslate"><span class="pre">lecun_normal()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.lecun_uniform"><code class="docutils literal notranslate"><span class="pre">lecun_uniform()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.normal"><code class="docutils literal notranslate"><span class="pre">normal()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.truncated_normal"><code class="docutils literal notranslate"><span class="pre">truncated_normal()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.ones"><code class="docutils literal notranslate"><span class="pre">ones()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.ones_init"><code class="docutils literal notranslate"><span class="pre">ones_init()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.orthogonal"><code class="docutils literal notranslate"><span class="pre">orthogonal()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.uniform"><code class="docutils literal notranslate"><span class="pre">uniform()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.variance_scaling"><code class="docutils literal notranslate"><span class="pre">variance_scaling()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.xavier_normal"><code class="docutils literal notranslate"><span class="pre">xavier_normal()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.xavier_uniform"><code class="docutils literal notranslate"><span class="pre">xavier_uniform()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.zeros"><code class="docutils literal notranslate"><span class="pre">zeros()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="initializers.html#flax.linen.initializers.zeros_init"><code class="docutils literal notranslate"><span class="pre">zeros_init()</span></code></a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="transformations.html">Transformations</a><ul> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.vmap"><code class="docutils literal notranslate"><span class="pre">vmap()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.scan"><code class="docutils literal notranslate"><span class="pre">scan()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.jit"><code class="docutils literal notranslate"><span class="pre">jit()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.remat"><code class="docutils literal notranslate"><span class="pre">remat()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.remat_scan"><code class="docutils literal notranslate"><span class="pre">remat_scan()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.map_variables"><code class="docutils literal notranslate"><span class="pre">map_variables()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.jvp"><code class="docutils literal notranslate"><span class="pre">jvp()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.vjp"><code class="docutils literal notranslate"><span class="pre">vjp()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.custom_vjp"><code class="docutils literal notranslate"><span class="pre">custom_vjp()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.while_loop"><code class="docutils literal notranslate"><span class="pre">while_loop()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.cond"><code class="docutils literal notranslate"><span class="pre">cond()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="transformations.html#flax.linen.switch"><code class="docutils literal notranslate"><span class="pre">switch()</span></code></a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="inspection.html">Inspection</a><ul> <li class="toctree-l2"><a class="reference internal" href="inspection.html#flax.linen.tabulate"><code class="docutils literal notranslate"><span class="pre">tabulate()</span></code></a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="variable.html">Variable dictionary</a><ul> <li class="toctree-l2"><a class="reference internal" href="variable.html#flax.linen.Variable"><code class="docutils literal notranslate"><span class="pre">Variable</span></code></a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="spmd.html">SPMD</a><ul> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.Partitioned"><code class="docutils literal notranslate"><span class="pre">Partitioned()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.with_partitioning"><code class="docutils literal notranslate"><span class="pre">with_partitioning()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.get_partition_spec"><code class="docutils literal notranslate"><span class="pre">get_partition_spec()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.get_sharding"><code class="docutils literal notranslate"><span class="pre">get_sharding()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.LogicallyPartitioned"><code class="docutils literal notranslate"><span class="pre">LogicallyPartitioned()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.logical_axis_rules"><code class="docutils literal notranslate"><span class="pre">logical_axis_rules()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.set_logical_axis_rules"><code class="docutils literal notranslate"><span class="pre">set_logical_axis_rules()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.get_logical_axis_rules"><code class="docutils literal notranslate"><span class="pre">get_logical_axis_rules()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.logical_to_mesh_axes"><code class="docutils literal notranslate"><span class="pre">logical_to_mesh_axes()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.logical_to_mesh"><code class="docutils literal notranslate"><span class="pre">logical_to_mesh()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.logical_to_mesh_sharding"><code class="docutils literal notranslate"><span class="pre">logical_to_mesh_sharding()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.with_logical_constraint"><code class="docutils literal notranslate"><span class="pre">with_logical_constraint()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="spmd.html#flax.linen.with_logical_partitioning"><code class="docutils literal notranslate"><span class="pre">with_logical_partitioning()</span></code></a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="decorators.html">Decorators</a><ul> <li class="toctree-l2"><a class="reference internal" href="decorators.html#flax.linen.compact"><code class="docutils literal notranslate"><span class="pre">compact()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="decorators.html#flax.linen.nowrap"><code class="docutils literal notranslate"><span class="pre">nowrap()</span></code></a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="profiling.html">Profiling</a><ul> <li class="toctree-l2"><a class="reference internal" href="profiling.html#flax.linen.enable_named_call"><code class="docutils literal notranslate"><span class="pre">enable_named_call()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="profiling.html#flax.linen.disable_named_call"><code class="docutils literal notranslate"><span class="pre">disable_named_call()</span></code></a></li> <li class="toctree-l2"><a class="reference internal" href="profiling.html#flax.linen.override_named_call"><code class="docutils literal notranslate"><span class="pre">override_named_call()</span></code></a></li> </ul> </li> </ul> </div> </div> </article> <footer class="prev-next-footer d-print-none"> <div class="prev-next-area"> <a class="left-prev" href="../flax.jax_utils.html" title="previous page"> <i class="fa-solid fa-angle-left"></i> <div class="prev-next-info"> <p class="prev-next-subtitle">previous</p> <p class="prev-next-title">flax.jax_utils package</p> </div> </a> <a class="right-next" href="module.html" title="next page"> <div class="prev-next-info"> <p class="prev-next-subtitle">next</p> <p class="prev-next-title">Module</p> </div> <i class="fa-solid fa-angle-right"></i> </a> </div> </footer> </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>