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>Module</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/module';</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="Init/Apply" href="init_apply.html" /> <link rel="prev" title="flax.linen" href="index.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/module.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"> 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"> <input type="search" class="form-control" name="q" id="search-input" placeholder="Search..." aria-label="Search..." autocomplete="off" autocorrect="off" autocapitalize="off" spellcheck="false"/> <kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd>K</kbd> </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. [Explore the new Flax NNX API ✨] </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"> Search <kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd> </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></summary><ul> <li class="toctree-l2 has-children"><a class="reference internal" href="../../guides/flax_fundamentals/index.html">Flax fundamentals</a><details><summary></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">setup</code> vs <code class="docutils literal notranslate">compact</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></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></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></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></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></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></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></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></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></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="reference internal" href="index.html">flax.linen</a><details open="open"><summary></summary><ul class="current"> <li class="toctree-l3 current active"><a class="current reference internal" href="#">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"> </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" > </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"> </button> <ul class="dropdown-menu"> <li><a href="../../_sources/api_reference/flax.linen/module.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" > .rst </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" > .pdf </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" > </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"> </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"> </button> `); </script> <button class="sidebar-toggle secondary-toggle btn btn-sm" title="Toggle secondary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip"> </button> </div></div> </div> </div> </div> <div id="jb-print-docs-body" class="onlyprint"> <h1>Module</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="#flax.linen.Module"><code class="docutils literal notranslate">Module</code></a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.__setattr__"><code class="docutils literal notranslate">Module.__setattr__()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.apply"><code class="docutils literal notranslate">Module.apply()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.bind"><code class="docutils literal notranslate">Module.bind()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.copy"><code class="docutils literal notranslate">Module.copy()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.get_variable"><code class="docutils literal notranslate">Module.get_variable()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.has_rng"><code class="docutils literal notranslate">Module.has_rng()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.has_variable"><code class="docutils literal notranslate">Module.has_variable()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.init"><code class="docutils literal notranslate">Module.init()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.init_with_output"><code class="docutils literal notranslate">Module.init_with_output()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.is_initializing"><code class="docutils literal notranslate">Module.is_initializing()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.is_mutable_collection"><code class="docutils literal notranslate">Module.is_mutable_collection()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.lazy_init"><code class="docutils literal notranslate">Module.lazy_init()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.make_rng"><code class="docutils literal notranslate">Module.make_rng()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.module_paths"><code class="docutils literal notranslate">Module.module_paths()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.param"><code class="docutils literal notranslate">Module.param()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.path"><code class="docutils literal notranslate">Module.path</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.perturb"><code class="docutils literal notranslate">Module.perturb()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.put_variable"><code class="docutils literal notranslate">Module.put_variable()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.setup"><code class="docutils literal notranslate">Module.setup()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.sow"><code class="docutils literal notranslate">Module.sow()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.tabulate"><code class="docutils literal notranslate">Module.tabulate()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.unbind"><code class="docutils literal notranslate">Module.unbind()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.variable"><code class="docutils literal notranslate">Module.variable()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.variables"><code class="docutils literal notranslate">Module.variables</code></a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.apply"><code class="docutils literal notranslate">apply()</code></a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.init"><code class="docutils literal notranslate">init()</code></a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.init_with_output"><code class="docutils literal notranslate">init_with_output()</code></a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.intercept_methods"><code class="docutils literal notranslate">intercept_methods()</code></a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.share_scope"><code class="docutils literal notranslate">share_scope()</code></a></li> </ul> </nav> </div> </div> </div> <div id="searchbox"></div> <article class="bd-article"> <div class="section" id="module-flax.linen"> <h1>Module<a class="headerlink" href="#module-flax.linen" title="Permalink to this heading">#</a></h1> The Flax Module system. <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.Module"> class flax.linen.Module<a class="reference internal" href="../../_modules/flax/linen/module.html#Module">[source]</a><a class="headerlink" href="#flax.linen.Module" title="Permalink to this definition">#</a></dt> <dd>Base class for all neural network modules. Layers and models should subclass this class. All Flax Modules are Python 3.7 <a class="reference external" href="https://docs.python.org/3/library/dataclasses.html">dataclasses</a>. Since dataclasses take over <code class="docutils literal notranslate">__init__</code>, you should instead override <a class="reference internal" href="#flax.linen.Module.setup" title="flax.linen.Module.setup"><code class="xref py py-meth docutils literal notranslate">setup()</code></a>, which is automatically called to initialize the module. Modules can contain submodules, and in this way can be nested in a tree structure. Submodels can be assigned as regular attributes inside the <a class="reference internal" href="#flax.linen.Module.setup" title="flax.linen.Module.setup"><code class="xref py py-meth docutils literal notranslate">setup()</code></a> method. You can define arbitrary “forward pass” methods on your Module subclass. While no methods are special-cased, <code class="docutils literal notranslate">__call__</code> is a popular choice because it allows you to use module instances as if they are functions: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> from flax import linen as nn >>> from typing import Tuple >>> class Module(nn.Module): ... features: Tuple[int, ...] = (16, 4) ... def setup(self): ... self.dense1 = nn.Dense(self.features[0]) ... self.dense2 = nn.Dense(self.features[1]) ... def __call__(self, x): ... return self.dense2(nn.relu(self.dense1(x))) </pre></div> </div> Optionally, for more concise module implementations where submodules definitions are co-located with their usage, you can use the <a class="reference internal" href="decorators.html#flax.linen.compact" title="flax.linen.compact"><code class="xref py py-meth docutils literal notranslate">compact()</code></a> wrapper. <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.__setattr__"> __setattr__(name, val)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.__setattr__">[source]</a><a class="headerlink" href="#flax.linen.Module.__setattr__" title="Permalink to this definition">#</a></dt> <dd>Sets an attribute on this Module. We overload setattr solely to support pythonic naming via assignment of submodules in the special <a class="reference internal" href="#flax.linen.Module.setup" title="flax.linen.Module.setup"><code class="xref py py-meth docutils literal notranslate">setup()</code></a> function: <div class="highlight-default notranslate"><div class="highlight"><pre>self.submodule_name = MyModule(...) </pre></div> </div> We also support lists and other general pytrees, e.g.: <div class="highlight-default notranslate"><div class="highlight"><pre>self.submodules = [MyModule0(..), MyModule1(..), ...] </pre></div> </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>name – Attribute to set.</li> <li>val – Value of the attribute.</li> </ul> </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.apply"> apply(variables, *args, rngs=None, method=None, mutable=False, capture_intermediates=False, **kwargs)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.apply">[source]</a><a class="headerlink" href="#flax.linen.Module.apply" title="Permalink to this definition">#</a></dt> <dd>Applies a module method to variables and returns output and modified variables. Note that <code class="docutils literal notranslate">method</code> should be set if one would like to call <code class="docutils literal notranslate">apply</code> on a different class method than <code class="docutils literal notranslate">__call__</code>. For instance, suppose a Transformer modules has a method called <code class="docutils literal notranslate">encode</code>, then the following calls <code class="docutils literal notranslate">apply</code> on that method: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> import numpy as np >>> class Transformer(nn.Module): ... def encode(self, x): ... ... >>> x = jnp.ones((16, 9)) >>> model = Transformer() >>> variables = model.init(jax.random.key(0), x, method=Transformer.encode) >>> encoded = model.apply(variables, x, method=Transformer.encode) </pre></div> </div> If a function instance is provided, the unbound function is used. For instance, the example below is equivalent to the one above: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> encoded = model.apply(variables, x, method=model.encode) </pre></div> </div> You can also pass a string to a callable attribute of the module. For example, the previous can be written as: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> encoded = model.apply(variables, x, method='encode') </pre></div> </div> Note <code class="docutils literal notranslate">method</code> can also be a function that is not defined in <code class="docutils literal notranslate">Transformer</code>. In that case, the function should have at least one argument representing an instance of the Module class: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> def other_fn(instance, x): ... # instance.some_module_attr(...) ... instance.encode ... ... >>> model.apply(variables, x, method=other_fn) </pre></div> </div> If you pass a single <code class="docutils literal notranslate">PRNGKey</code>, Flax will use it to feed the <code class="docutils literal notranslate">'params'</code> RNG stream. If you want to use a different RNG stream or need to use multiple streams, you can pass a dictionary mapping each RNG stream name to its corresponding <code class="docutils literal notranslate">PRNGKey</code> to <code class="docutils literal notranslate">apply</code>. If <code class="docutils literal notranslate">self.make_rng(name)</code> is called on an RNG stream name that isn’t passed by the user, it will default to using the <code class="docutils literal notranslate">'params'</code> RNG stream. Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x, add_noise=False): ... x = nn.Dense(16)(x) ... x = nn.relu(x) ... ... if add_noise: ... # Add gaussian noise ... noise_key = self.make_rng('noise') ... x = x + jax.random.normal(noise_key, x.shape) ... ... return nn.Dense(1)(x) >>> x = jnp.empty((1, 7)) >>> module = Foo() >>> rngs = {'params': jax.random.key(0), 'noise': jax.random.key(1)} >>> variables = module.init(rngs, x) >>> out0 = module.apply(variables, x, add_noise=True, rngs=rngs) >>> rngs['noise'] = jax.random.key(0) >>> out1 = module.apply(variables, x, add_noise=True, rngs=rngs) >>> # different output (key(1) vs key(0)) >>> np.testing.assert_raises(AssertionError, np.testing.assert_allclose, out0, out1) >>> del rngs['noise'] >>> # self.make_rng('noise') will default to using the 'params' RNG stream >>> out2 = module.apply(variables, x, add_noise=True, rngs=rngs) >>> # same output (key(0)) >>> np.testing.assert_allclose(out1, out2) >>> # passing in a single key is equivalent to passing in {'params': key} >>> out3 = module.apply(variables, x, add_noise=True, rngs=jax.random.key(0)) >>> # same output (key(0)) >>> np.testing.assert_allclose(out2, out3) </pre></div> </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>variables – A dictionary containing variables keyed by variable collections. See <a class="reference internal" href="variable.html#module-flax.core.variables" title="flax.core.variables"><code class="xref py py-mod docutils literal notranslate">flax.core.variables</code></a> for more details about variables.</li> <li>*args – Named arguments passed to the specified apply method.</li> <li>rngs – a dict of PRNGKeys to initialize the PRNG sequences. The “params” PRNG sequence is used to initialize parameters.</li> <li>method – A function to call apply on. This is generally a function in the module. If provided, applies this method. If not provided, applies the <code class="docutils literal notranslate">__call__</code> method of the module. A string can also be provided to specify a method by name.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections.</li> <li>capture_intermediates – If <code class="docutils literal notranslate">True</code>, captures intermediate return values of all Modules inside the “intermediates” collection. By default, only the return values of all <code class="docutils literal notranslate">__call__</code> methods are stored. A function can be passed to change the filter behavior. The filter function takes the Module instance and method name and returns a bool indicating whether the output of that method invocation should be stored.</li> <li>**kwargs – Keyword arguments passed to the specified apply method.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">If <code class="docutils literal notranslate">mutable</code> is False, returns output. If any collections are mutable, returns <code class="docutils literal notranslate">(output, vars)</code>, where <code class="docutils literal notranslate">vars</code> are is a dict of the modified collections. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.bind"> bind(variables, *args, rngs=None, mutable=False)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.bind">[source]</a><a class="headerlink" href="#flax.linen.Module.bind" title="Permalink to this definition">#</a></dt> <dd>Creates an interactive Module instance by binding variables and RNGs. <code class="docutils literal notranslate">bind</code> provides an “interactive” instance of a Module directly without transforming a function with <code class="docutils literal notranslate">apply</code>. This is particularly useful for debugging and interactive use cases like notebooks where a function would limit the ability to split up code into different cells. Once the variables (and optionally RNGs) are bound to a <code class="docutils literal notranslate">Module</code> it becomes a stateful object. Note that idiomatic JAX is functional and therefore an interactive instance does not mix well with vanilla JAX APIs. <code class="docutils literal notranslate">bind()</code> should only be used for interactive experimentation, and in all other cases we strongly encourage users to use <code class="docutils literal notranslate">apply()</code> instead. Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import jax >>> import jax.numpy as jnp >>> import flax.linen as nn >>> class AutoEncoder(nn.Module): ... def setup(self): ... self.encoder = nn.Dense(3) ... self.decoder = nn.Dense(5) ... ... def __call__(self, x): ... return self.decoder(self.encoder(x)) >>> x = jnp.ones((16, 9)) >>> ae = AutoEncoder() >>> variables = ae.init(jax.random.key(0), x) >>> model = ae.bind(variables) >>> z = model.encoder(x) >>> x_reconstructed = model.decoder(z) </pre></div> </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>variables – A dictionary containing variables keyed by variable collections. See <a class="reference internal" href="variable.html#module-flax.core.variables" title="flax.core.variables"><code class="xref py py-mod docutils literal notranslate">flax.core.variables</code></a> for more details about variables.</li> <li>*args – Named arguments (not used).</li> <li>rngs – a dict of PRNGKeys to initialize the PRNG sequences.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A copy of this instance with bound variables and RNGs. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.copy"> copy(*, parent=<flax.linen.module._Sentinel object>, name=None, **updates)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.copy">[source]</a><a class="headerlink" href="#flax.linen.Module.copy" title="Permalink to this definition">#</a></dt> <dd>Creates a copy of this Module, with optionally updated arguments. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>parent – The parent of the copy. By default the current module is taken as parent if not explicitly specified.</li> <li>name – A new name for the copied Module, by default a new automatic name will be given.</li> <li>**updates – Attribute updates.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A copy of the this Module with the updated name, parent, and attributes. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.get_variable"> get_variable(col, name, default=None)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.get_variable">[source]</a><a class="headerlink" href="#flax.linen.Module.get_variable" title="Permalink to this definition">#</a></dt> <dd>Retrieves the value of a Variable. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>col – the variable collection.</li> <li>name – the name of the variable.</li> <li>default – the default value to return if the variable does not exist in this scope.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The value of the input variable, of the default value if the variable doesn’t exist in this scope. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.has_rng"> has_rng(name)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.has_rng">[source]</a><a class="headerlink" href="#flax.linen.Module.has_rng" title="Permalink to this definition">#</a></dt> <dd>Returns true if a PRNGSequence with name <code class="docutils literal notranslate">name</code> exists. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.has_variable"> has_variable(col, name)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.has_variable">[source]</a><a class="headerlink" href="#flax.linen.Module.has_variable" title="Permalink to this definition">#</a></dt> <dd>Checks if a variable of given collection and name exists in this Module. See <a class="reference internal" href="variable.html#module-flax.core.variables" title="flax.core.variables"><code class="xref py py-mod docutils literal notranslate">flax.core.variables</code></a> for more explanation on variables and collections. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>col – The variable collection name.</li> <li>name – The name of the variable.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">True if the variable exists. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.init"> init(rngs, *args, method=None, mutable=DenyList(deny='intermediates'), capture_intermediates=False, **kwargs)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.init">[source]</a><a class="headerlink" href="#flax.linen.Module.init" title="Permalink to this definition">#</a></dt> <dd>Initializes a module method with variables and returns modified variables. <code class="docutils literal notranslate">init</code> takes as first argument either a single <code class="docutils literal notranslate">PRNGKey</code>, or a dictionary mapping variable collections names to their <code class="docutils literal notranslate">PRNGKeys</code>, and will call <code class="docutils literal notranslate">method</code> (which is the module’s <code class="docutils literal notranslate">__call__</code> function by default) passing <code class="docutils literal notranslate">*args</code> and <code class="docutils literal notranslate">**kwargs</code>, and returns a dictionary of initialized variables. Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> import numpy as np >>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x, train): ... x = nn.Dense(16)(x) ... x = nn.BatchNorm(use_running_average=not train)(x) ... x = nn.relu(x) ... return nn.Dense(1)(x) >>> x = jnp.empty((1, 7)) >>> module = Foo() >>> key = jax.random.key(0) >>> variables = module.init(key, x, train=False) </pre></div> </div> If you pass a single <code class="docutils literal notranslate">PRNGKey</code>, Flax will use it to feed the <code class="docutils literal notranslate">'params'</code> RNG stream. If you want to use a different RNG stream or need to use multiple streams, you can pass a dictionary mapping each RNG stream name to its corresponding <code class="docutils literal notranslate">PRNGKey</code> to <code class="docutils literal notranslate">init</code>. If <code class="docutils literal notranslate">self.make_rng(name)</code> is called on an RNG stream name that isn’t passed by the user, it will default to using the <code class="docutils literal notranslate">'params'</code> RNG stream. Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x): ... x = nn.Dense(16)(x) ... x = nn.relu(x) ... ... other_variable = self.variable( ... 'other_collection', ... 'other_variable', ... lambda x: jax.random.normal(self.make_rng('other_rng'), x.shape), ... x, ... ) ... x = x + other_variable.value ... ... return nn.Dense(1)(x) >>> module = Foo() >>> rngs = {'params': jax.random.key(0), 'other_rng': jax.random.key(1)} >>> variables0 = module.init(rngs, x) >>> rngs['other_rng'] = jax.random.key(0) >>> variables1 = module.init(rngs, x) >>> # equivalent params (key(0)) >>> _ = jax.tree_util.tree_map( ... np.testing.assert_allclose, variables0['params'], variables1['params'] ... ) >>> # different other_variable (key(1) vs key(0)) >>> np.testing.assert_raises( ... AssertionError, ... np.testing.assert_allclose, ... variables0['other_collection']['other_variable'], ... variables1['other_collection']['other_variable'], ... ) >>> del rngs['other_rng'] >>> # self.make_rng('other_rng') will default to using the 'params' RNG stream >>> variables2 = module.init(rngs, x) >>> # equivalent params (key(0)) >>> _ = jax.tree_util.tree_map( ... np.testing.assert_allclose, variables1['params'], variables2['params'] ... ) >>> # equivalent other_variable (key(0)) >>> np.testing.assert_allclose( ... variables1['other_collection']['other_variable'], ... variables2['other_collection']['other_variable'], ... ) >>> # passing in a single key is equivalent to passing in {'params': key} >>> variables3 = module.init(jax.random.key(0), x) >>> # equivalent params (key(0)) >>> _ = jax.tree_util.tree_map( ... np.testing.assert_allclose, variables2['params'], variables3['params'] ... ) >>> # equivalent other_variable (key(0)) >>> np.testing.assert_allclose( ... variables2['other_collection']['other_variable'], ... variables3['other_collection']['other_variable'], ... ) </pre></div> </div> Jitting <code class="docutils literal notranslate">init</code> initializes a model lazily using only the shapes of the provided arguments, and avoids computing the forward pass with actual values. Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> module = nn.Dense(1) >>> init_jit = jax.jit(module.init) >>> variables = init_jit(jax.random.key(0), x) </pre></div> </div> <code class="docutils literal notranslate">init</code> is a light wrapper over <code class="docutils literal notranslate">apply</code>, so other <code class="docutils literal notranslate">apply</code> arguments like <code class="docutils literal notranslate">method</code>, <code class="docutils literal notranslate">mutable</code>, and <code class="docutils literal notranslate">capture_intermediates</code> are also available. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rngs – The rngs for the variable collections.</li> <li>*args – Named arguments passed to the init function.</li> <li>method – An optional method. If provided, applies this method. If not provided, applies the <code class="docutils literal notranslate">__call__</code> method. A string can also be provided to specify a method by name.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections. By default all collections except “intermediates” are mutable.</li> <li>capture_intermediates – If <code class="docutils literal notranslate">True</code>, captures intermediate return values of all Modules inside the “intermediates” collection. By default only the return values of all <code class="docutils literal notranslate">__call__</code> methods are stored. A function can be passed to change the filter behavior. The filter function takes the Module instance and method name and returns a bool indicating whether the output of that method invocation should be stored.</li> <li>**kwargs – Keyword arguments passed to the init function.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The initialized variable dict. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.init_with_output"> init_with_output(rngs, *args, method=None, mutable=DenyList(deny='intermediates'), capture_intermediates=False, **kwargs)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.init_with_output">[source]</a><a class="headerlink" href="#flax.linen.Module.init_with_output" title="Permalink to this definition">#</a></dt> <dd>Initializes a module method with variables and returns output and modified variables. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rngs – The rngs for the variable collections.</li> <li>*args – Named arguments passed to the init function.</li> <li>method – An optional method. If provided, applies this method. If not provided, applies the <code class="docutils literal notranslate">__call__</code> method. A string can also be provided to specify a method by name.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections. By default, all collections except “intermediates” are mutable.</li> <li>capture_intermediates – If <code class="docutils literal notranslate">True</code>, captures intermediate return values of all Modules inside the “intermediates” collection. By default only the return values of all <code class="docutils literal notranslate">__call__</code> methods are stored. A function can be passed to change the filter behavior. The filter function takes the Module instance and method name and returns a bool indicating whether the output of that method invocation should be stored.</li> <li>**kwargs – Keyword arguments passed to the init function.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even"><code class="docutils literal notranslate">(output, vars)</code>, where <code class="docutils literal notranslate">vars</code> are is a dict of the modified collections. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.is_initializing"> is_initializing()<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.is_initializing">[source]</a><a class="headerlink" href="#flax.linen.Module.is_initializing" title="Permalink to this definition">#</a></dt> <dd>Returns True if running under self.init(…) or nn.init(…)(). This is a helper method to handle the common case of simple initialization where we wish to have setup logic occur when only called under <code class="docutils literal notranslate">module.init</code> or <code class="docutils literal notranslate">nn.init</code>. For more complicated multi-phase initialization scenarios it is better to test for the mutability of particular variable collections or for the presence of particular variables that potentially need to be initialized. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.is_mutable_collection"> is_mutable_collection(col)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.is_mutable_collection">[source]</a><a class="headerlink" href="#flax.linen.Module.is_mutable_collection" title="Permalink to this definition">#</a></dt> <dd>Returns true if the collection <code class="docutils literal notranslate">col</code> is mutable. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.lazy_init"> lazy_init(rngs, *args, method=None, mutable=DenyList(deny='intermediates'), **kwargs)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.lazy_init">[source]</a><a class="headerlink" href="#flax.linen.Module.lazy_init" title="Permalink to this definition">#</a></dt> <dd>Initializes a module without computing on an actual input. lazy_init will initialize the variables without doing unnecessary compute. The input data should be passed as a <code class="docutils literal notranslate">jax.ShapeDtypeStruct</code> which specifies the shape and dtype of the input but no concrete data. Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> model = nn.Dense(features=256) >>> variables = model.lazy_init( ... jax.random.key(0), jax.ShapeDtypeStruct((1, 128), jnp.float32)) </pre></div> </div> The args and kwargs args passed to <code class="docutils literal notranslate">lazy_init</code> can be a mix of concrete (jax arrays, scalars, bools) and abstract (ShapeDtypeStruct) values. Concrete values are only necessary for arguments that affect the initialization of variables. For example, the model might expect a keyword arg that enables/disables a subpart of the model. In this case, an explicit value (True/Flase) should be passed otherwise <code class="docutils literal notranslate">lazy_init</code> cannot infer which variables should be initialized. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rngs – The rngs for the variable collections.</li> <li>*args – arguments passed to the init function.</li> <li>method – An optional method. If provided, applies this method. If not provided, applies the <code class="docutils literal notranslate">__call__</code> method.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections. By default all collections except “intermediates” are mutable.</li> <li>**kwargs – Keyword arguments passed to the init function.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The initialized variable dict. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.make_rng"> make_rng(name='params')<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.make_rng">[source]</a><a class="headerlink" href="#flax.linen.Module.make_rng" title="Permalink to this definition">#</a></dt> <dd>Returns a new RNG key from a given RNG sequence for this Module. The new RNG key is split from the previous one. Thus, every call to <code class="docutils literal notranslate">make_rng</code> returns a new RNG key, while still guaranteeing full reproducibility. <div class="admonition note"> Note If an invalid name is passed (i.e. no RNG key was passed by the user in <code class="docutils literal notranslate">.init</code> or <code class="docutils literal notranslate">.apply</code> for this name), then <code class="docutils literal notranslate">name</code> will default to <code class="docutils literal notranslate">'params'</code>. </div> Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import jax >>> import flax.linen as nn >>> class ParamsModule(nn.Module): ... def __call__(self): ... return self.make_rng('params') >>> class OtherModule(nn.Module): ... def __call__(self): ... return self.make_rng('other') >>> key = jax.random.key(0) >>> params_out, _ = ParamsModule().init_with_output({'params': key}) >>> # self.make_rng('other') will default to using the 'params' RNG stream >>> other_out, _ = OtherModule().init_with_output({'params': key}) >>> assert params_out == other_out </pre></div> </div> Learn more about RNG’s by reading the Flax RNG guide: <a class="reference external" href="https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/rng_guide.html">https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/rng_guide.html</a> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd">name – The RNG sequence name. </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The newly generated RNG key. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.module_paths"> module_paths(rngs, *args, show_repeated=False, mutable=DenyList(deny='intermediates'), **kwargs)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.module_paths">[source]</a><a class="headerlink" href="#flax.linen.Module.module_paths" title="Permalink to this definition">#</a></dt> <dd>Returns a dictionary mapping module paths to module instances. This method has the same signature and internally calls <code class="docutils literal notranslate">Module.init</code>, but instead of returning the variables, it returns a dictionary mapping module paths to unbounded copies of module instances that were used at runtime. <code class="docutils literal notranslate">module_paths</code> uses <code class="docutils literal notranslate">jax.eval_shape</code> to run the forward computation without consuming any FLOPs or allocating memory. Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x): ... h = nn.Dense(4)(x) ... return nn.Dense(2)(h) >>> x = jnp.ones((16, 9)) >>> modules = Foo().module_paths(jax.random.key(0), x) >>> print({ ... p: type(m).__name__ for p, m in modules.items() ... }) {'': 'Foo', 'Dense_0': 'Dense', 'Dense_1': 'Dense'} </pre></div> </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rngs – The rngs for the variable collections as passed to <code class="docutils literal notranslate">Module.init</code>.</li> <li>*args – The arguments to the forward computation.</li> <li>show_repeated – If <code class="docutils literal notranslate">True</code>, repeated calls to the same module will be shown in the table, otherwise only the first call will be shown. Default is <code class="docutils literal notranslate">False</code>.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections. By default, all collections except ‘intermediates’ are mutable.</li> <li>**kwargs – keyword arguments to pass to the forward computation.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A dict`ionary mapping module paths to module instances. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.param"> param(name, init_fn, *init_args, unbox=True, **init_kwargs)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.param">[source]</a><a class="headerlink" href="#flax.linen.Module.param" title="Permalink to this definition">#</a></dt> <dd>Declares and returns a parameter in this Module. Parameters are read-only variables in the collection named “params”. See <a class="reference internal" href="variable.html#module-flax.core.variables" title="flax.core.variables"><code class="xref py py-mod docutils literal notranslate">flax.core.variables</code></a> for more details on variables. The first argument of <code class="docutils literal notranslate">init_fn</code> is assumed to be a PRNG key, which is provided automatically and does not have to be passed using <code class="docutils literal notranslate">init_args</code> or <code class="docutils literal notranslate">init_kwargs</code>: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x): ... x = nn.Dense(4)(x) ... mean = self.param('mean', nn.initializers.lecun_normal(), x.shape) ... ... ... return x * mean >>> variables = Foo().init({'params': jax.random.key(0), 'stats': jax.random.key(1)}, jnp.ones((2, 3))) >>> jax.tree_util.tree_map(jnp.shape, variables) {'params': {'Dense_0': {'bias': (4,), 'kernel': (3, 4)}, 'mean': (2, 4)}} </pre></div> </div> In the example above, the function <code class="docutils literal notranslate">lecun_normal</code> expects two arguments: <code class="docutils literal notranslate">key</code> and <code class="docutils literal notranslate">shape</code>, but only <code class="docutils literal notranslate">shape</code> has to be provided explicitly; <code class="docutils literal notranslate">key</code> is set automatically using the PRNG for <code class="docutils literal notranslate">params</code> that is passed when initializing the module using <a class="reference internal" href="#flax.linen.init" title="flax.linen.init"><code class="xref py py-meth docutils literal notranslate">init()</code></a>. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>name – The parameter name.</li> <li>init_fn – The function that will be called to compute the initial value of this variable. This function will only be called the first time this parameter is used in this module.</li> <li>*init_args – The positional arguments to pass to init_fn.</li> <li>unbox – If True, <code class="docutils literal notranslate">AxisMetadata</code> instances are replaced by their unboxed value, see <code class="docutils literal notranslate">flax.nn.meta.unbox</code> (default: True).</li> <li>**init_kwargs – The key-word arguments to pass to init_fn.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The value of the initialized parameter. Throws an error if the parameter exists already. </dd> </dl> </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="flax.linen.Module.path"> property path<a class="headerlink" href="#flax.linen.Module.path" title="Permalink to this definition">#</a></dt> <dd>Get the path of this Module. Top-level root modules have an empty path <code class="docutils literal notranslate">()</code>. Note that this method can only be used on bound modules that have a valid scope. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> class SubModel(nn.Module): ... @nn.compact ... def __call__(self, x): ... print(f'SubModel path: {self.path}') ... return x >>> class Model(nn.Module): ... @nn.compact ... def __call__(self, x): ... print(f'Model path: {self.path}') ... return SubModel()(x) >>> model = Model() >>> variables = model.init(jax.random.key(0), jnp.ones((1, 2))) Model path: () SubModel path: ('SubModel_0',) </pre></div> </div> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.perturb"> perturb(name, value, collection='perturbations')<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.perturb">[source]</a><a class="headerlink" href="#flax.linen.Module.perturb" title="Permalink to this definition">#</a></dt> <dd>Add an zero-value variable (‘perturbation’) to the intermediate value. The gradient of <code class="docutils literal notranslate">value</code> would be the same as the gradient of this perturbation variable. Therefore, if you define your loss function with both params and perturbations as standalone arguments, you can get the intermediate gradients of <code class="docutils literal notranslate">value</code> by running <code class="docutils literal notranslate">jax.grad</code> on the perturbation argument. <div class="admonition note"> Note This is an experimental API and may be tweaked later for better performance and usability. At its current stage, it creates extra dummy variables that occupies extra memory space. Use it only to debug gradients in training. </div> Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x): ... x = nn.Dense(3)(x) ... x = self.perturb('dense3', x) ... return nn.Dense(2)(x) >>> def loss(variables, inputs, targets): ... preds = model.apply(variables, inputs) ... return jnp.square(preds - targets).mean() >>> x = jnp.ones((2, 9)) >>> y = jnp.ones((2, 2)) >>> model = Foo() >>> variables = model.init(jax.random.key(0), x) >>> intm_grads = jax.grad(loss, argnums=0)(variables, x, y) >>> print(intm_grads['perturbations']['dense3']) [[-1.456924 -0.44332537 0.02422847] [-1.456924 -0.44332537 0.02422847]] </pre></div> </div> If perturbations are not passed to <code class="docutils literal notranslate">apply</code>, <code class="docutils literal notranslate">perturb</code> behaves like a no-op so you can easily disable the behavior when not needed: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> model.apply(variables, x) # works as expected Array([[-1.0980128 , -0.67961735], [-1.0980128 , -0.67961735]], dtype=float32) >>> model.apply({'params': variables['params']}, x) # behaves like a no-op Array([[-1.0980128 , -0.67961735], [-1.0980128 , -0.67961735]], dtype=float32) >>> intm_grads = jax.grad(loss, argnums=0)({'params': variables['params']}, x, y) >>> 'perturbations' not in intm_grads True </pre></div> </div> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.put_variable"> put_variable(col, name, value)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.put_variable">[source]</a><a class="headerlink" href="#flax.linen.Module.put_variable" title="Permalink to this definition">#</a></dt> <dd>Updates the value of the given variable if it is mutable, or an error otherwise. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>col – the variable collection.</li> <li>name – the name of the variable.</li> <li>value – the new value of the variable.</li> </ul> </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.setup"> setup()<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.setup">[source]</a><a class="headerlink" href="#flax.linen.Module.setup" title="Permalink to this definition">#</a></dt> <dd>Initializes a Module lazily (similar to a lazy <code class="docutils literal notranslate">__init__</code>). <code class="docutils literal notranslate">setup</code> is called once lazily on a module instance when a module is bound, immediately before any other methods like <code class="docutils literal notranslate">__call__</code> are invoked, or before a <code class="docutils literal notranslate">setup</code>-defined attribute on <code class="docutils literal notranslate">self</code> is accessed. This can happen in three cases: <blockquote> <div><ol class="arabic"> <li>Immediately when invoking <a class="reference internal" href="#flax.linen.apply" title="flax.linen.apply"><code class="xref py py-meth docutils literal notranslate">apply()</code></a>, <a class="reference internal" href="#flax.linen.init" title="flax.linen.init"><code class="xref py py-meth docutils literal notranslate">init()</code></a> or <code class="xref py py-meth docutils literal notranslate">init_and_output()</code>.</li> <li>Once the module is given a name by being assigned to an attribute of another module inside the other module’s <code class="docutils literal notranslate">setup</code> method (see <a class="reference internal" href="#flax.linen.Module.__setattr__" title="flax.linen.Module.__setattr__"><code class="xref py py-meth docutils literal notranslate">__setattr__()</code></a>): <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class MyModule(nn.Module): ... def setup(self): ... submodule = nn.Conv(...) ... # Accessing `submodule` attributes does not yet work here. ... # The following line invokes `self.__setattr__`, which gives ... # `submodule` the name "conv1". ... self.conv1 = submodule ... # Accessing `submodule` attributes or methods is now safe and ... # either causes setup() to be called once. </pre></div> </div> </li> <li>Once a module is constructed inside a method wrapped with <a class="reference internal" href="decorators.html#flax.linen.compact" title="flax.linen.compact"><code class="xref py py-meth docutils literal notranslate">compact()</code></a>, immediately before another method is called or <code class="docutils literal notranslate">setup</code> defined attribute is accessed.</li> </ol> </div></blockquote> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.sow"> sow(col, name, value, reduce_fn=<function <lambda>>, init_fn=<function <lambda>>)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.sow">[source]</a><a class="headerlink" href="#flax.linen.Module.sow" title="Permalink to this definition">#</a></dt> <dd>Stores a value in a collection. Collections can be used to collect intermediate values without the overhead of explicitly passing a container through each Module call. If the target collection is not mutable <code class="docutils literal notranslate">sow</code> behaves like a no-op and returns <code class="docutils literal notranslate">False</code>. Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import jax >>> import jax.numpy as jnp >>> import flax.linen as nn >>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x): ... h = nn.Dense(4)(x) ... self.sow('intermediates', 'h', h) ... return nn.Dense(2)(h) >>> x = jnp.ones((16, 9)) >>> model = Foo() >>> variables = model.init(jax.random.key(0), x) >>> y, state = model.apply(variables, x, mutable=['intermediates']) >>> jax.tree.map(jnp.shape, state['intermediates']) {'h': ((16, 4),)} </pre></div> </div> By default the values are stored in a tuple and each stored value is appended at the end. This way all intermediates can be tracked when the same module is called multiple times. Alternatively, a custom init/reduce function can be passed: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Foo2(nn.Module): ... @nn.compact ... def __call__(self, x): ... init_fn = lambda: 0 ... reduce_fn = lambda a, b: a + b ... self.sow('intermediates', 'h', x, ... init_fn=init_fn, reduce_fn=reduce_fn) ... self.sow('intermediates', 'h', x * 2, ... init_fn=init_fn, reduce_fn=reduce_fn) ... return x >>> x = jnp.ones((1, 1)) >>> model = Foo2() >>> variables = model.init(jax.random.key(0), x) >>> y, state = model.apply( ... variables, x, mutable=['intermediates']) >>> print(state['intermediates']) {'h': Array([[3.]], dtype=float32)} </pre></div> </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>col – The name of the variable collection.</li> <li>name – The name of the variable.</li> <li>value – The value of the variable.</li> <li>reduce_fn – The function used to combine the existing value with the new value. The default is to append the value to a tuple.</li> <li>init_fn – For the first value stored, <code class="docutils literal notranslate">reduce_fn</code> will be passed the result of <code class="docutils literal notranslate">init_fn</code> together with the value to be stored. The default is an empty tuple.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even"><code class="docutils literal notranslate">True</code> if the value has been stored successfully, <code class="docutils literal notranslate">False</code> otherwise. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.tabulate"> tabulate(rngs, *args, depth=None, show_repeated=False, mutable=DenyList(deny='intermediates'), console_kwargs=None, table_kwargs=mappingproxy({}), column_kwargs=mappingproxy({}), compute_flops=False, compute_vjp_flops=False, **kwargs)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.tabulate">[source]</a><a class="headerlink" href="#flax.linen.Module.tabulate" title="Permalink to this definition">#</a></dt> <dd>Creates a summary of the Module represented as a table. This method has the same signature and internally calls <code class="docutils literal notranslate">Module.init</code>, but instead of returning the variables, it returns the string summarizing the Module in a table. <code class="docutils literal notranslate">tabulate</code> uses <code class="docutils literal notranslate">jax.eval_shape</code> to run the forward computation without consuming any FLOPs or allocating memory. Additional arguments can be passed into the <code class="docutils literal notranslate">console_kwargs</code> argument, for example, <code class="docutils literal notranslate">{'width': 120}</code>. For a full list of <code class="docutils literal notranslate">console_kwargs</code> arguments, see: <a class="reference external" href="https://rich.readthedocs.io/en/stable/reference/console.html#rich.console.Console">https://rich.readthedocs.io/en/stable/reference/console.html#rich.console.Console</a> Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x): ... h = nn.Dense(4)(x) ... return nn.Dense(2)(h) >>> x = jnp.ones((16, 9)) >>> # print(Foo().tabulate( >>> # jax.random.key(0), x, compute_flops=True, compute_vjp_flops=True)) </pre></div> </div> This gives the following output: <div class="highlight-default notranslate"><div class="highlight"><pre> Foo Summary ┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓ ┃ path ┃ module ┃ inputs ┃ outputs ┃ flops ┃ vjp_flops ┃ params ┃ ┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩ │ │ Foo │ float32[16,9] │ float32[16,2] │ 1504 │ 4460 │ │ ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤ │ Dense_0 │ Dense │ float32[16,9] │ float32[16,4] │ 1216 │ 3620 │ bias: │ │ │ │ │ │ │ │ float32[4] │ │ │ │ │ │ │ │ kernel: │ │ │ │ │ │ │ │ float32[9,4] │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ 40 (160 B) │ ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤ │ Dense_1 │ Dense │ float32[16,4] │ float32[16,2] │ 288 │ 840 │ bias: │ │ │ │ │ │ │ │ float32[2] │ │ │ │ │ │ │ │ kernel: │ │ │ │ │ │ │ │ float32[4,2] │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ 10 (40 B) │ ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤ │ │ │ │ │ │ Total │ 50 (200 B) │ └─────────┴────────┴───────────────┴───────────────┴───────┴───────────┴─────────────────┘ Total Parameters: 50 (200 B) </pre></div> </div> Note: rows order in the table does not represent execution order, instead it aligns with the order of keys in <code class="docutils literal notranslate">variables</code> which are sorted alphabetically. Note: <code class="docutils literal notranslate">vjp_flops</code> returns <code class="docutils literal notranslate">0</code> if the module is not differentiable. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rngs – The rngs for the variable collections as passed to <code class="docutils literal notranslate">Module.init</code>.</li> <li>*args – The arguments to the forward computation.</li> <li>depth – controls how many submodule deep the summary can go. By default, its <code class="docutils literal notranslate">None</code> which means no limit. If a submodule is not shown because of the depth limit, its parameter count and bytes will be added to the row of its first shown ancestor such that the sum of all rows always adds up to the total number of parameters of the Module.</li> <li>show_repeated – If <code class="docutils literal notranslate">True</code>, repeated calls to the same module will be shown in the table, otherwise only the first call will be shown. Default is <code class="docutils literal notranslate">False</code>.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections. By default, all collections except ‘intermediates’ are mutable.</li> <li>console_kwargs – An optional dictionary with additional keyword arguments that are passed to <code class="docutils literal notranslate">rich.console.Console</code> when rendering the table. Default arguments are <code class="docutils literal notranslate">{'force_terminal': True, 'force_jupyter': False}</code>.</li> <li>table_kwargs – An optional dictionary with additional keyword arguments that are passed to <code class="docutils literal notranslate">rich.table.Table</code> constructor.</li> <li>column_kwargs – An optional dictionary with additional keyword arguments that are passed to <code class="docutils literal notranslate">rich.table.Table.add_column</code> when adding columns to the table.</li> <li>compute_flops – whether to include a <code class="docutils literal notranslate">flops</code> column in the table listing the estimated FLOPs cost of each module forward pass. Does incur actual on-device computation / compilation / memory allocation, but still introduces overhead for large modules (e.g. extra 20 seconds for a Stable Diffusion’s UNet, whereas otherwise tabulation would finish in 5 seconds).</li> <li>compute_vjp_flops – whether to include a <code class="docutils literal notranslate">vjp_flops</code> column in the table listing the estimated FLOPs cost of each module backward pass. Introduces a compute overhead of about 2-3X of <code class="docutils literal notranslate">compute_flops</code>.</li> <li>**kwargs – keyword arguments to pass to the forward computation.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A string summarizing the Module. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.unbind"> unbind()<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.unbind">[source]</a><a class="headerlink" href="#flax.linen.Module.unbind" title="Permalink to this definition">#</a></dt> <dd>Returns an unbound copy of a Module and its variables. <code class="docutils literal notranslate">unbind</code> helps create a stateless version of a bound Module. An example of a common use case: to extract a sub-Module defined inside <code class="docutils literal notranslate">setup()</code> and its corresponding variables: 1) temporarily <code class="docutils literal notranslate">bind</code> the parent Module; and then 2) <code class="docutils literal notranslate">unbind</code> the desired sub-Module. (Recall that <code class="docutils literal notranslate">setup()</code> is only called when the Module is bound.): <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Encoder(nn.Module): ... @nn.compact ... def __call__(self, x): ... ... ... return nn.Dense(256)(x) >>> class Decoder(nn.Module): ... @nn.compact ... def __call__(self, x): ... ... ... return nn.Dense(784)(x) >>> class AutoEncoder(nn.Module): ... def setup(self): ... self.encoder = Encoder() ... self.decoder = Decoder() ... ... def __call__(self, x): ... return self.decoder(self.encoder(x)) >>> module = AutoEncoder() >>> variables = module.init(jax.random.key(0), jnp.ones((1, 784))) >>> # Extract the Encoder sub-Module and its variables >>> encoder, encoder_vars = module.bind(variables).encoder.unbind() </pre></div> </div> <dl class="field-list simple"> <dt class="field-odd">Returns</dt> <dd class="field-odd">A tuple with an unbound copy of this Module and its variables. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Module.variable"> variable(col, name, init_fn=None, *init_args, unbox=True, **init_kwargs)<a class="reference internal" href="../../_modules/flax/linen/module.html#Module.variable">[source]</a><a class="headerlink" href="#flax.linen.Module.variable" title="Permalink to this definition">#</a></dt> <dd>Declares and returns a variable in this Module. See <a class="reference internal" href="variable.html#module-flax.core.variables" title="flax.core.variables"><code class="xref py py-mod docutils literal notranslate">flax.core.variables</code></a> for more information. See also <a class="reference internal" href="#flax.linen.Module.param" title="flax.linen.Module.param"><code class="xref py py-meth docutils literal notranslate">param()</code></a> for a shorthand way to define read-only variables in the “params” collection. Contrary to <a class="reference internal" href="#flax.linen.Module.param" title="flax.linen.Module.param"><code class="xref py py-meth docutils literal notranslate">param()</code></a>, all arguments passing using <code class="docutils literal notranslate">init_fn</code> should be passed on explicitly: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x): ... x = nn.Dense(4)(x) ... key = self.make_rng('stats') ... mean = self.variable('stats', 'mean', nn.initializers.lecun_normal(), key, x.shape) ... ... ... return x * mean.value >>> variables = Foo().init({'params': jax.random.key(0), 'stats': jax.random.key(1)}, jnp.ones((2, 3))) >>> jax.tree_util.tree_map(jnp.shape, variables) {'params': {'Dense_0': {'bias': (4,), 'kernel': (3, 4)}}, 'stats': {'mean': (2, 4)}} </pre></div> </div> In the example above, the function <code class="docutils literal notranslate">lecun_normal</code> expects two arguments: <code class="docutils literal notranslate">key</code> and <code class="docutils literal notranslate">shape</code>, and both have to be passed on. The PRNG for <code class="docutils literal notranslate">stats</code> has to be provided explicitly when calling <a class="reference internal" href="#flax.linen.init" title="flax.linen.init"><code class="xref py py-meth docutils literal notranslate">init()</code></a> and <a class="reference internal" href="#flax.linen.apply" title="flax.linen.apply"><code class="xref py py-meth docutils literal notranslate">apply()</code></a>. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>col – The variable collection name.</li> <li>name – The variable name.</li> <li>init_fn – The function that will be called to compute the initial value of this variable. This function will only be called the first time this variable is used in this module. If None, the variable must already be initialized otherwise an error is raised.</li> <li>*init_args – The positional arguments to pass to init_fn.</li> <li>unbox – If True, <code class="docutils literal notranslate">AxisMetadata</code> instances are replaced by their unboxed value, see <code class="docutils literal notranslate">flax.nn.meta.unbox</code> (default: True).</li> <li>**init_kwargs – The key-word arguments to pass to init_fn</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A <code class="xref py py-class docutils literal notranslate">flax.core.variables.Variable</code> that can be read or set via “.value” attribute. Throws an error if the variable exists already. </dd> </dl> </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="flax.linen.Module.variables"> property variables<a class="headerlink" href="#flax.linen.Module.variables" title="Permalink to this definition">#</a></dt> <dd>Returns the variables in this module. </dd></dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.apply"> flax.linen.apply(fn, module, mutable=False, capture_intermediates=False)<a class="reference internal" href="../../_modules/flax/linen/module.html#apply">[source]</a><a class="headerlink" href="#flax.linen.apply" title="Permalink to this definition">#</a></dt> <dd>Creates an apply function to call <code class="docutils literal notranslate">fn</code> with a bound module. Unlike <code class="docutils literal notranslate">Module.apply</code> this function returns a new function with the signature <code class="docutils literal notranslate">(variables, *args, rngs=None, **kwargs) -> T</code> where <code class="docutils literal notranslate">T</code> is the return type of <code class="docutils literal notranslate">fn</code>. If <code class="docutils literal notranslate">mutable</code> is not <code class="docutils literal notranslate">False</code> the return type is a tuple where the second item is a <code class="docutils literal notranslate">FrozenDict</code> with the mutated variables. The apply function that is returned can be directly composed with JAX transformations like <code class="docutils literal notranslate">jax.jit</code>: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Foo(nn.Module): ... def encode(self, x): ... ... ... def decode(self, x): ... ... >>> def f(foo, x): ... z = foo.encode(x) ... y = foo.decode(z) ... # ... ... return y >>> variables = {} >>> foo = Foo() >>> f_jitted = jax.jit(nn.apply(f, foo)) >>> f_jitted(variables, jnp.ones((1, 3))) </pre></div> </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>fn – The function that should be applied. The first argument passed will be a module instance of the <code class="docutils literal notranslate">module</code> with variables and RNGs bound to it.</li> <li>module – The <code class="docutils literal notranslate">Module</code> that will be used to bind variables and RNGs to. The <code class="docutils literal notranslate">Module</code> passed as the first argument to <code class="docutils literal notranslate">fn</code> will be a clone of module.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections.</li> <li>capture_intermediates – If <code class="docutils literal notranslate">True</code>, captures intermediate return values of all Modules inside the “intermediates” collection. By default, only the return values of all <cite>__call__</cite> methods are stored. A function can be passed to change the filter behavior. The filter function takes the Module instance and method name and returns a bool indicating whether the output of that method invocation should be stored.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The apply function wrapping <code class="docutils literal notranslate">fn</code>. </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.init"> flax.linen.init(fn, module, mutable=DenyList(deny='intermediates'), capture_intermediates=False)<a class="reference internal" href="../../_modules/flax/linen/module.html#init">[source]</a><a class="headerlink" href="#flax.linen.init" title="Permalink to this definition">#</a></dt> <dd>Creates an init function to call <code class="docutils literal notranslate">fn</code> with a bound module. Unlike <code class="docutils literal notranslate">Module.init</code> this function returns a new function with the signature <code class="docutils literal notranslate">(rngs, *args, **kwargs) -> variables</code>. The rngs can be a dict of PRNGKeys or a single <code class="docutils literal notranslate">`PRNGKey</code> which is equivalent to passing a dict with one PRNGKey with the name “params”. The init function that is returned can be directly composed with JAX transformations like <code class="docutils literal notranslate">jax.jit</code>: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Foo(nn.Module): ... def encode(self, x): ... ... ... def decode(self, x): ... ... >>> def f(foo, x): ... z = foo.encode(x) ... y = foo.decode(z) ... # ... ... return y >>> foo = Foo() >>> f_jitted = jax.jit(nn.init(f, foo)) >>> variables = f_jitted(jax.random.key(0), jnp.ones((1, 3))) </pre></div> </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>fn – The function that should be applied. The first argument passed will be a module instance of the <code class="docutils literal notranslate">module</code> with variables and RNGs bound to it.</li> <li>module – The <code class="docutils literal notranslate">Module</code> that will be used to bind variables and RNGs to. The <code class="docutils literal notranslate">Module</code> passed as the first argument to <code class="docutils literal notranslate">fn</code> will be a clone of module.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections. By default, all collections except “intermediates” are mutable.</li> <li>capture_intermediates – If <cite>True</cite>, captures intermediate return values of all Modules inside the “intermediates” collection. By default, only the return values of all <cite>__call__</cite> methods are stored. A function can be passed to change the filter behavior. The filter function takes the Module instance and method name and returns a bool indicating whether the output of that method invocation should be stored.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The init function wrapping <code class="docutils literal notranslate">fn</code>. </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.init_with_output"> flax.linen.init_with_output(fn, module, mutable=DenyList(deny='intermediates'), capture_intermediates=False)<a class="reference internal" href="../../_modules/flax/linen/module.html#init_with_output">[source]</a><a class="headerlink" href="#flax.linen.init_with_output" title="Permalink to this definition">#</a></dt> <dd>Creates an init function to call <code class="docutils literal notranslate">fn</code> with a bound module that also returns the function outputs. Unlike <code class="docutils literal notranslate">Module.init_with_output</code> this function returns a new function with the signature <code class="docutils literal notranslate">(rngs, *args, **kwargs) -> (T, variables)</code> where <code class="docutils literal notranslate">T</code> is the return type of <code class="docutils literal notranslate">fn</code>. The rngs can be a dict of PRNGKeys or a single <code class="docutils literal notranslate">`PRNGKey</code> which is equivalent to passing a dict with one PRNGKey with the name “params”. The init function that is returned can be directly composed with JAX transformations like <code class="docutils literal notranslate">jax.jit</code>: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class Foo(nn.Module): ... def encode(self, x): ... ... ... def decode(self, x): ... ... >>> def f(foo, x): ... z = foo.encode(x) ... y = foo.decode(z) ... # ... ... return y >>> foo = Foo() >>> f_jitted = jax.jit(nn.init_with_output(f, foo)) >>> y, variables = f_jitted(jax.random.key(0), jnp.ones((1, 3))) </pre></div> </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>fn – The function that should be applied. The first argument passed will be a module instance of the <code class="docutils literal notranslate">module</code> with variables and RNGs bound to it.</li> <li>module – The <code class="docutils literal notranslate">Module</code> that will be used to bind variables and RNGs to. The <code class="docutils literal notranslate">Module</code> passed as the first argument to <code class="docutils literal notranslate">fn</code> will be a clone of module.</li> <li>mutable – Can be bool, str, or list. Specifies which collections should be treated as mutable: <code class="docutils literal notranslate">bool</code>: all/no collections are mutable. <code class="docutils literal notranslate">str</code>: The name of a single mutable collection. <code class="docutils literal notranslate">list</code>: A list of names of mutable collections. By default, all collections except “intermediates” are mutable.</li> <li>capture_intermediates – If <code class="docutils literal notranslate">True</code>, captures intermediate return values of all Modules inside the “intermediates” collection. By default, only the return values of all <cite>__call__</cite> methods are stored. A function can be passed to change the filter behavior. The filter function takes the Module instance and method name and returns a bool indicating whether the output of that method invocation should be stored.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The init function wrapping <code class="docutils literal notranslate">fn</code>. </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.intercept_methods"> flax.linen.intercept_methods(interceptor)<a class="reference internal" href="../../_modules/flax/linen/module.html#intercept_methods">[source]</a><a class="headerlink" href="#flax.linen.intercept_methods" title="Permalink to this definition">#</a></dt> <dd>Registers a new method interceptor. Method interceptors allow you to (at a distance) intercept method calls to modules. It works similarly to decorators. You could modify args/kwargs before calling the underlying method and/or modify the result returning from calling the underlying method. Or you could completely skip calling the underlying method and decide to do something differently. For example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax.numpy as jnp ... >>> class Foo(nn.Module): ... def __call__(self, x): ... return x ... >>> def my_interceptor1(next_fun, args, kwargs, context): ... print('calling my_interceptor1') ... return next_fun(*args, **kwargs) ... >>> foo = Foo() >>> with nn.intercept_methods(my_interceptor1): ... _ = foo(jnp.ones([1])) calling my_interceptor1 </pre></div> </div> You could also register multiple interceptors on the same method. Interceptors will run in order. For example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> def my_interceptor2(next_fun, args, kwargs, context): ... print('calling my_interceptor2') ... return next_fun(*args, **kwargs) ... >>> with nn.intercept_methods(my_interceptor1), \ ... nn.intercept_methods(my_interceptor2): ... _ = foo(jnp.ones([1])) calling my_interceptor1 calling my_interceptor2 </pre></div> </div> You could skip other interceptors by directly calling the <code class="docutils literal notranslate">context.orig_method</code>. For example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> def my_interceptor3(next_fun, args, kwargs, context): ... print('calling my_interceptor3') ... return context.orig_method(*args, **kwargs) >>> with nn.intercept_methods(my_interceptor3), \ ... nn.intercept_methods(my_interceptor1), \ ... nn.intercept_methods(my_interceptor2): ... _ = foo(jnp.ones([1])) calling my_interceptor3 </pre></div> </div> The following methods couldn’t be intercepted: <ol class="arabic simple"> <li>Methods decoratored with <code class="docutils literal notranslate">nn.nowrap</code>.</li> <li>Dunder methods including <code class="docutils literal notranslate">__eq__</code>, <code class="docutils literal notranslate">__repr__</code>, <code class="docutils literal notranslate">__init__</code>, <code class="docutils literal notranslate">__hash__</code>, and <code class="docutils literal notranslate">__post_init__</code>.</li> <li>Module dataclass fields.</li> <li>Module descriptors.</li> </ol> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd">interceptor – A method interceptor. </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.share_scope"> flax.linen.share_scope(module, other, /)<a class="reference internal" href="../../_modules/flax/linen/module.html#share_scope">[source]</a><a class="headerlink" href="#flax.linen.share_scope" title="Permalink to this definition">#</a></dt> <dd>Modifies one of the Modules such that they share the same scope. This is useful when you want to wrap a Module and extend its functionality without changing the parameter structure. <code class="docutils literal notranslate">share_scope</code> takes two Modules, <code class="docutils literal notranslate">module</code> and <code class="docutils literal notranslate">other</code>. <code class="docutils literal notranslate">module</code> will use <code class="docutils literal notranslate">other</code>’s scope if <code class="docutils literal notranslate">other</code> has a scope and its not a descendant of``module``’s scope: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax >>> from jax import numpy as jnp, random ... >>> class DenseLoRA(nn.Module): ... base: nn.Dense ... rank: int ... ... def setup(self): ... nn.share_scope(self, self.base) ... ... @nn.compact ... def __call__(self, x: jax.Array): ... din, dout = x.shape[-1], self.base.features ... A = self.param('A', nn.zeros_init(), (din, self.rank)) ... B = self.param('B', nn.zeros_init(), (self.rank, dout)) ... return self.base(x) + x @ A @ B ... >>> class Model(nn.Module): ... @nn.compact ... def __call__(self, x: jax.Array): ... dense = nn.Dense(10) # base scope ... return DenseLoRA(dense, rank=2)(x) # reuse the base scope ... >>> model = Model() ... >>> params = model.init(random.key(0), jnp.ones((1, 5)))['params'] >>> list(params['Dense_0'].keys()) ['A', 'B', 'kernel', 'bias'] </pre></div> </div> When <code class="docutils literal notranslate">other</code>’s scope is a descendant of <code class="docutils literal notranslate">module</code>’s scope then <code class="docutils literal notranslate">other</code> will use <code class="docutils literal notranslate">module</code>’s scope instead: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class DenseLoRA(nn.Module): ... features: int ... rank: int ... ... def setup(self): ... self.child = nn.Dense(self.features) ... nn.share_scope(self, self.child) ... ... @nn.compact ... def __call__(self, x: jax.Array): ... din, dout = x.shape[-1], self.features ... A = self.param('A', nn.zeros_init(), (din, self.rank)) ... B = self.param('B', nn.zeros_init(), (self.rank, dout)) ... return self.child(x) + x @ A @ B ... >>> class Model(nn.Module): ... @nn.compact ... def __call__(self, x: jax.Array): ... return DenseLoRA(10, rank=2)(x) ... >>> model = Model() ... >>> params = model.init(random.key(0), jnp.ones((1, 5)))['params'] >>> list(params['DenseLoRA_0'].keys()) ['A', 'B', 'kernel', 'bias'] </pre></div> </div> </dd></dl> </div> </article> <footer class="prev-next-footer d-print-none"> <div class="prev-next-area"> <a class="left-prev" href="index.html" title="previous page"> <div class="prev-next-info"> previous flax.linen </div> </a> <a class="right-next" href="init_apply.html" title="next page"> <div class="prev-next-info"> next Init/Apply </div> </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"> 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="#flax.linen.Module"><code class="docutils literal notranslate">Module</code></a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.__setattr__"><code class="docutils literal notranslate">Module.__setattr__()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.apply"><code class="docutils literal notranslate">Module.apply()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.bind"><code class="docutils literal notranslate">Module.bind()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.copy"><code class="docutils literal notranslate">Module.copy()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.get_variable"><code class="docutils literal notranslate">Module.get_variable()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.has_rng"><code class="docutils literal notranslate">Module.has_rng()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.has_variable"><code class="docutils literal notranslate">Module.has_variable()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.init"><code class="docutils literal notranslate">Module.init()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.init_with_output"><code class="docutils literal notranslate">Module.init_with_output()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.is_initializing"><code class="docutils literal notranslate">Module.is_initializing()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.is_mutable_collection"><code class="docutils literal notranslate">Module.is_mutable_collection()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.lazy_init"><code class="docutils literal notranslate">Module.lazy_init()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.make_rng"><code class="docutils literal notranslate">Module.make_rng()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.module_paths"><code class="docutils literal notranslate">Module.module_paths()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.param"><code class="docutils literal notranslate">Module.param()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.path"><code class="docutils literal notranslate">Module.path</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.perturb"><code class="docutils literal notranslate">Module.perturb()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.put_variable"><code class="docutils literal notranslate">Module.put_variable()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.setup"><code class="docutils literal notranslate">Module.setup()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.sow"><code class="docutils literal notranslate">Module.sow()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.tabulate"><code class="docutils literal notranslate">Module.tabulate()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.unbind"><code class="docutils literal notranslate">Module.unbind()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.variable"><code class="docutils literal notranslate">Module.variable()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Module.variables"><code class="docutils literal notranslate">Module.variables</code></a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.apply"><code class="docutils literal notranslate">apply()</code></a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.init"><code class="docutils literal notranslate">init()</code></a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.init_with_output"><code class="docutils literal notranslate">init_with_output()</code></a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.intercept_methods"><code class="docutils literal notranslate">intercept_methods()</code></a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.share_scope"><code class="docutils literal notranslate">share_scope()</code></a></li> </ul> </nav></div> </div></div> </div> <footer class="bd-footer-content"> <div class="bd-footer-content__inner container"> <div class="footer-item"> By The Flax authors </div> <div class="footer-item"> © Copyright 2023, The Flax authors. </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>