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>Layers</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>window.MathJax = {"options": {"processHtmlClass": "tex2jax_process|mathjax_process|math|output_area"}}</script> <script defer="defer" src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script> <script>DOCUMENTATION_OPTIONS.pagename = 'api_reference/flax.linen/layers';</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="Activation functions" href="activation_functions.html" /> <link rel="prev" title="Init/Apply" href="init_apply.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/layers.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"><a class="reference internal" href="module.html">Module</a></li> <li class="toctree-l3"><a class="reference internal" href="init_apply.html">Init/Apply</a></li> <li class="toctree-l3 current active"><a class="current reference internal" href="#">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/layers.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>Layers</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="#linear-modules">Linear Modules</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.Dense"><code class="docutils literal notranslate">Dense</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.features"><code class="docutils literal notranslate">Dense.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.use_bias"><code class="docutils literal notranslate">Dense.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.dtype"><code class="docutils literal notranslate">Dense.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.param_dtype"><code class="docutils literal notranslate">Dense.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.precision"><code class="docutils literal notranslate">Dense.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.kernel_init"><code class="docutils literal notranslate">Dense.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.bias_init"><code class="docutils literal notranslate">Dense.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.__call__"><code class="docutils literal notranslate">Dense.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral"><code class="docutils literal notranslate">DenseGeneral</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.features"><code class="docutils literal notranslate">DenseGeneral.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.axis"><code class="docutils literal notranslate">DenseGeneral.axis</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.batch_dims"><code class="docutils literal notranslate">DenseGeneral.batch_dims</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.use_bias"><code class="docutils literal notranslate">DenseGeneral.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.dtype"><code class="docutils literal notranslate">DenseGeneral.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.param_dtype"><code class="docutils literal notranslate">DenseGeneral.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.kernel_init"><code class="docutils literal notranslate">DenseGeneral.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.bias_init"><code class="docutils literal notranslate">DenseGeneral.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.precision"><code class="docutils literal notranslate">DenseGeneral.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.__call__"><code class="docutils literal notranslate">DenseGeneral.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv"><code class="docutils literal notranslate">Conv</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.features"><code class="docutils literal notranslate">Conv.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.kernel_size"><code class="docutils literal notranslate">Conv.kernel_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.strides"><code class="docutils literal notranslate">Conv.strides</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.padding"><code class="docutils literal notranslate">Conv.padding</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.input_dilation"><code class="docutils literal notranslate">Conv.input_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.kernel_dilation"><code class="docutils literal notranslate">Conv.kernel_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.feature_group_count"><code class="docutils literal notranslate">Conv.feature_group_count</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.use_bias"><code class="docutils literal notranslate">Conv.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.mask"><code class="docutils literal notranslate">Conv.mask</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.dtype"><code class="docutils literal notranslate">Conv.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.param_dtype"><code class="docutils literal notranslate">Conv.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.precision"><code class="docutils literal notranslate">Conv.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.kernel_init"><code class="docutils literal notranslate">Conv.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.bias_init"><code class="docutils literal notranslate">Conv.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.__call__"><code class="docutils literal notranslate">Conv.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose"><code class="docutils literal notranslate">ConvTranspose</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.features"><code class="docutils literal notranslate">ConvTranspose.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.kernel_size"><code class="docutils literal notranslate">ConvTranspose.kernel_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.strides"><code class="docutils literal notranslate">ConvTranspose.strides</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.padding"><code class="docutils literal notranslate">ConvTranspose.padding</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.kernel_dilation"><code class="docutils literal notranslate">ConvTranspose.kernel_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.use_bias"><code class="docutils literal notranslate">ConvTranspose.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.mask"><code class="docutils literal notranslate">ConvTranspose.mask</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.dtype"><code class="docutils literal notranslate">ConvTranspose.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.param_dtype"><code class="docutils literal notranslate">ConvTranspose.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.precision"><code class="docutils literal notranslate">ConvTranspose.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.kernel_init"><code class="docutils literal notranslate">ConvTranspose.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.bias_init"><code class="docutils literal notranslate">ConvTranspose.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.transpose_kernel"><code class="docutils literal notranslate">ConvTranspose.transpose_kernel</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.__call__"><code class="docutils literal notranslate">ConvTranspose.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal"><code class="docutils literal notranslate">ConvLocal</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.features"><code class="docutils literal notranslate">ConvLocal.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.kernel_size"><code class="docutils literal notranslate">ConvLocal.kernel_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.strides"><code class="docutils literal notranslate">ConvLocal.strides</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.padding"><code class="docutils literal notranslate">ConvLocal.padding</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.input_dilation"><code class="docutils literal notranslate">ConvLocal.input_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.kernel_dilation"><code class="docutils literal notranslate">ConvLocal.kernel_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.feature_group_count"><code class="docutils literal notranslate">ConvLocal.feature_group_count</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.use_bias"><code class="docutils literal notranslate">ConvLocal.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.mask"><code class="docutils literal notranslate">ConvLocal.mask</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.dtype"><code class="docutils literal notranslate">ConvLocal.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.param_dtype"><code class="docutils literal notranslate">ConvLocal.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.precision"><code class="docutils literal notranslate">ConvLocal.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.kernel_init"><code class="docutils literal notranslate">ConvLocal.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.bias_init"><code class="docutils literal notranslate">ConvLocal.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.__call__"><code class="docutils literal notranslate">ConvLocal.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum"><code class="docutils literal notranslate">Einsum</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.shape"><code class="docutils literal notranslate">Einsum.shape</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.einsum_str"><code class="docutils literal notranslate">Einsum.einsum_str</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.use_bias"><code class="docutils literal notranslate">Einsum.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.dtype"><code class="docutils literal notranslate">Einsum.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.param_dtype"><code class="docutils literal notranslate">Einsum.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.precision"><code class="docutils literal notranslate">Einsum.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.kernel_init"><code class="docutils literal notranslate">Einsum.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.bias_init"><code class="docutils literal notranslate">Einsum.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.__call__"><code class="docutils literal notranslate">Einsum.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed"><code class="docutils literal notranslate">Embed</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.num_embeddings"><code class="docutils literal notranslate">Embed.num_embeddings</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.features"><code class="docutils literal notranslate">Embed.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.dtype"><code class="docutils literal notranslate">Embed.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.param_dtype"><code class="docutils literal notranslate">Embed.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.embedding_init"><code class="docutils literal notranslate">Embed.embedding_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.__call__"><code class="docutils literal notranslate">Embed.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.attend"><code class="docutils literal notranslate">Embed.attend()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#pooling">Pooling</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.max_pool"><code class="docutils literal notranslate">max_pool()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.avg_pool"><code class="docutils literal notranslate">avg_pool()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.pool"><code class="docutils literal notranslate">pool()</code></a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#normalization">Normalization</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.BatchNorm"><code class="docutils literal notranslate">BatchNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.use_running_average"><code class="docutils literal notranslate">BatchNorm.use_running_average</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.axis"><code class="docutils literal notranslate">BatchNorm.axis</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.momentum"><code class="docutils literal notranslate">BatchNorm.momentum</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.epsilon"><code class="docutils literal notranslate">BatchNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.dtype"><code class="docutils literal notranslate">BatchNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.param_dtype"><code class="docutils literal notranslate">BatchNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.use_bias"><code class="docutils literal notranslate">BatchNorm.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.use_scale"><code class="docutils literal notranslate">BatchNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.bias_init"><code class="docutils literal notranslate">BatchNorm.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.scale_init"><code class="docutils literal notranslate">BatchNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.axis_name"><code class="docutils literal notranslate">BatchNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.axis_index_groups"><code class="docutils literal notranslate">BatchNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.use_fast_variance"><code class="docutils literal notranslate">BatchNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.__call__"><code class="docutils literal notranslate">BatchNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm"><code class="docutils literal notranslate">LayerNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.epsilon"><code class="docutils literal notranslate">LayerNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.dtype"><code class="docutils literal notranslate">LayerNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.param_dtype"><code class="docutils literal notranslate">LayerNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.use_bias"><code class="docutils literal notranslate">LayerNorm.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.use_scale"><code class="docutils literal notranslate">LayerNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.bias_init"><code class="docutils literal notranslate">LayerNorm.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.scale_init"><code class="docutils literal notranslate">LayerNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.reduction_axes"><code class="docutils literal notranslate">LayerNorm.reduction_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.feature_axes"><code class="docutils literal notranslate">LayerNorm.feature_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.axis_name"><code class="docutils literal notranslate">LayerNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.axis_index_groups"><code class="docutils literal notranslate">LayerNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.use_fast_variance"><code class="docutils literal notranslate">LayerNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.__call__"><code class="docutils literal notranslate">LayerNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm"><code class="docutils literal notranslate">GroupNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.num_groups"><code class="docutils literal notranslate">GroupNorm.num_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.group_size"><code class="docutils literal notranslate">GroupNorm.group_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.epsilon"><code class="docutils literal notranslate">GroupNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.dtype"><code class="docutils literal notranslate">GroupNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.param_dtype"><code class="docutils literal notranslate">GroupNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.use_bias"><code class="docutils literal notranslate">GroupNorm.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.use_scale"><code class="docutils literal notranslate">GroupNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.bias_init"><code class="docutils literal notranslate">GroupNorm.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.scale_init"><code class="docutils literal notranslate">GroupNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.reduction_axes"><code class="docutils literal notranslate">GroupNorm.reduction_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.axis_name"><code class="docutils literal notranslate">GroupNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.axis_index_groups"><code class="docutils literal notranslate">GroupNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.use_fast_variance"><code class="docutils literal notranslate">GroupNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.__call__"><code class="docutils literal notranslate">GroupNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm"><code class="docutils literal notranslate">RMSNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.epsilon"><code class="docutils literal notranslate">RMSNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.dtype"><code class="docutils literal notranslate">RMSNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.param_dtype"><code class="docutils literal notranslate">RMSNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.use_scale"><code class="docutils literal notranslate">RMSNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.scale_init"><code class="docutils literal notranslate">RMSNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.reduction_axes"><code class="docutils literal notranslate">RMSNorm.reduction_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.feature_axes"><code class="docutils literal notranslate">RMSNorm.feature_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.axis_name"><code class="docutils literal notranslate">RMSNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.axis_index_groups"><code class="docutils literal notranslate">RMSNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.use_fast_variance"><code class="docutils literal notranslate">RMSNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.__call__"><code class="docutils literal notranslate">RMSNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm"><code class="docutils literal notranslate">InstanceNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.epsilon"><code class="docutils literal notranslate">InstanceNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.dtype"><code class="docutils literal notranslate">InstanceNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.param_dtype"><code class="docutils literal notranslate">InstanceNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.use_bias"><code class="docutils literal notranslate">InstanceNorm.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.use_scale"><code class="docutils literal notranslate">InstanceNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.bias_init"><code class="docutils literal notranslate">InstanceNorm.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.scale_init"><code class="docutils literal notranslate">InstanceNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.feature_axes"><code class="docutils literal notranslate">InstanceNorm.feature_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.axis_name"><code class="docutils literal notranslate">InstanceNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.axis_index_groups"><code class="docutils literal notranslate">InstanceNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.use_fast_variance"><code class="docutils literal notranslate">InstanceNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.__call__"><code class="docutils literal notranslate">InstanceNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm"><code class="docutils literal notranslate">SpectralNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.layer_instance"><code class="docutils literal notranslate">SpectralNorm.layer_instance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.n_steps"><code class="docutils literal notranslate">SpectralNorm.n_steps</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.epsilon"><code class="docutils literal notranslate">SpectralNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.dtype"><code class="docutils literal notranslate">SpectralNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.param_dtype"><code class="docutils literal notranslate">SpectralNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.error_on_non_matrix"><code class="docutils literal notranslate">SpectralNorm.error_on_non_matrix</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.collection_name"><code class="docutils literal notranslate">SpectralNorm.collection_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.__call__"><code class="docutils literal notranslate">SpectralNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm"><code class="docutils literal notranslate">WeightNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.layer_instance"><code class="docutils literal notranslate">WeightNorm.layer_instance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.epsilon"><code class="docutils literal notranslate">WeightNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.dtype"><code class="docutils literal notranslate">WeightNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.param_dtype"><code class="docutils literal notranslate">WeightNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.use_scale"><code class="docutils literal notranslate">WeightNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.scale_init"><code class="docutils literal notranslate">WeightNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.feature_axes"><code class="docutils literal notranslate">WeightNorm.feature_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.variable_filter"><code class="docutils literal notranslate">WeightNorm.variable_filter</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.__call__"><code class="docutils literal notranslate">WeightNorm.__call__()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#combinators">Combinators</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.Sequential"><code class="docutils literal notranslate">Sequential</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Sequential.layers"><code class="docutils literal notranslate">Sequential.layers</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Sequential.__call__"><code class="docutils literal notranslate">Sequential.__call__()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#stochastic">Stochastic</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.Dropout"><code class="docutils literal notranslate">Dropout</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.rate"><code class="docutils literal notranslate">Dropout.rate</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.broadcast_dims"><code class="docutils literal notranslate">Dropout.broadcast_dims</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.deterministic"><code class="docutils literal notranslate">Dropout.deterministic</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.rng_collection"><code class="docutils literal notranslate">Dropout.rng_collection</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.__call__"><code class="docutils literal notranslate">Dropout.__call__()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#attention">Attention</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.MultiHeadDotProductAttention"><code class="docutils literal notranslate">MultiHeadDotProductAttention</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.num_heads"><code class="docutils literal notranslate">MultiHeadDotProductAttention.num_heads</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.dtype"><code class="docutils literal notranslate">MultiHeadDotProductAttention.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.param_dtype"><code class="docutils literal notranslate">MultiHeadDotProductAttention.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.qkv_features"><code class="docutils literal notranslate">MultiHeadDotProductAttention.qkv_features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.out_features"><code class="docutils literal notranslate">MultiHeadDotProductAttention.out_features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.broadcast_dropout"><code class="docutils literal notranslate">MultiHeadDotProductAttention.broadcast_dropout</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.dropout_rate"><code class="docutils literal notranslate">MultiHeadDotProductAttention.dropout_rate</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.deterministic"><code class="docutils literal notranslate">MultiHeadDotProductAttention.deterministic</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.precision"><code class="docutils literal notranslate">MultiHeadDotProductAttention.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.kernel_init"><code class="docutils literal notranslate">MultiHeadDotProductAttention.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.out_kernel_init"><code class="docutils literal notranslate">MultiHeadDotProductAttention.out_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.bias_init"><code class="docutils literal notranslate">MultiHeadDotProductAttention.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.out_bias_init"><code class="docutils literal notranslate">MultiHeadDotProductAttention.out_bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.use_bias"><code class="docutils literal notranslate">MultiHeadDotProductAttention.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.attention_fn"><code class="docutils literal notranslate">MultiHeadDotProductAttention.attention_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.decode"><code class="docutils literal notranslate">MultiHeadDotProductAttention.decode</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.normalize_qk"><code class="docutils literal notranslate">MultiHeadDotProductAttention.normalize_qk</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.qk_attn_weights_einsum_cls"><code class="docutils literal notranslate">MultiHeadDotProductAttention.qk_attn_weights_einsum_cls</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.attn_weights_value_einsum_cls"><code class="docutils literal notranslate">MultiHeadDotProductAttention.attn_weights_value_einsum_cls</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.__call__"><code class="docutils literal notranslate">MultiHeadDotProductAttention.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention"><code class="docutils literal notranslate">MultiHeadAttention</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.num_heads"><code class="docutils literal notranslate">MultiHeadAttention.num_heads</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.dtype"><code class="docutils literal notranslate">MultiHeadAttention.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.param_dtype"><code class="docutils literal notranslate">MultiHeadAttention.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.qkv_features"><code class="docutils literal notranslate">MultiHeadAttention.qkv_features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.out_features"><code class="docutils literal notranslate">MultiHeadAttention.out_features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.broadcast_dropout"><code class="docutils literal notranslate">MultiHeadAttention.broadcast_dropout</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.dropout_rate"><code class="docutils literal notranslate">MultiHeadAttention.dropout_rate</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.deterministic"><code class="docutils literal notranslate">MultiHeadAttention.deterministic</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.precision"><code class="docutils literal notranslate">MultiHeadAttention.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.kernel_init"><code class="docutils literal notranslate">MultiHeadAttention.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.bias_init"><code class="docutils literal notranslate">MultiHeadAttention.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.use_bias"><code class="docutils literal notranslate">MultiHeadAttention.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.attention_fn"><code class="docutils literal notranslate">MultiHeadAttention.attention_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.decode"><code class="docutils literal notranslate">MultiHeadAttention.decode</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.normalize_qk"><code class="docutils literal notranslate">MultiHeadAttention.normalize_qk</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.__call__"><code class="docutils literal notranslate">MultiHeadAttention.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SelfAttention"><code class="docutils literal notranslate">SelfAttention</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SelfAttention.__call__"><code class="docutils literal notranslate">SelfAttention.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.dot_product_attention_weights"><code class="docutils literal notranslate">dot_product_attention_weights()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.dot_product_attention"><code class="docutils literal notranslate">dot_product_attention()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.make_attention_mask"><code class="docutils literal notranslate">make_attention_mask()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.make_causal_mask"><code class="docutils literal notranslate">make_causal_mask()</code></a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#recurrent">Recurrent</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.RNNCellBase"><code class="docutils literal notranslate">RNNCellBase</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNNCellBase.__call__"><code class="docutils literal notranslate">RNNCellBase.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNNCellBase.initialize_carry"><code class="docutils literal notranslate">RNNCellBase.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell"><code class="docutils literal notranslate">LSTMCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.features"><code class="docutils literal notranslate">LSTMCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.gate_fn"><code class="docutils literal notranslate">LSTMCell.gate_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.activation_fn"><code class="docutils literal notranslate">LSTMCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.kernel_init"><code class="docutils literal notranslate">LSTMCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.recurrent_kernel_init"><code class="docutils literal notranslate">LSTMCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.bias_init"><code class="docutils literal notranslate">LSTMCell.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.dtype"><code class="docutils literal notranslate">LSTMCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.param_dtype"><code class="docutils literal notranslate">LSTMCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.__call__"><code class="docutils literal notranslate">LSTMCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.initialize_carry"><code class="docutils literal notranslate">LSTMCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell"><code class="docutils literal notranslate">OptimizedLSTMCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.gate_fn"><code class="docutils literal notranslate">OptimizedLSTMCell.gate_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.activation_fn"><code class="docutils literal notranslate">OptimizedLSTMCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.kernel_init"><code class="docutils literal notranslate">OptimizedLSTMCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.recurrent_kernel_init"><code class="docutils literal notranslate">OptimizedLSTMCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.bias_init"><code class="docutils literal notranslate">OptimizedLSTMCell.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.dtype"><code class="docutils literal notranslate">OptimizedLSTMCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.param_dtype"><code class="docutils literal notranslate">OptimizedLSTMCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.__call__"><code class="docutils literal notranslate">OptimizedLSTMCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.initialize_carry"><code class="docutils literal notranslate">OptimizedLSTMCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell"><code class="docutils literal notranslate">ConvLSTMCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.features"><code class="docutils literal notranslate">ConvLSTMCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.kernel_size"><code class="docutils literal notranslate">ConvLSTMCell.kernel_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.strides"><code class="docutils literal notranslate">ConvLSTMCell.strides</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.padding"><code class="docutils literal notranslate">ConvLSTMCell.padding</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.bias"><code class="docutils literal notranslate">ConvLSTMCell.bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.dtype"><code class="docutils literal notranslate">ConvLSTMCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.param_dtype"><code class="docutils literal notranslate">ConvLSTMCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.__call__"><code class="docutils literal notranslate">ConvLSTMCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.initialize_carry"><code class="docutils literal notranslate">ConvLSTMCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell"><code class="docutils literal notranslate">SimpleCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.features"><code class="docutils literal notranslate">SimpleCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.activation_fn"><code class="docutils literal notranslate">SimpleCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.kernel_init"><code class="docutils literal notranslate">SimpleCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.recurrent_kernel_init"><code class="docutils literal notranslate">SimpleCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.bias_init"><code class="docutils literal notranslate">SimpleCell.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.dtype"><code class="docutils literal notranslate">SimpleCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.param_dtype"><code class="docutils literal notranslate">SimpleCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.residual"><code class="docutils literal notranslate">SimpleCell.residual</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.__call__"><code class="docutils literal notranslate">SimpleCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.initialize_carry"><code class="docutils literal notranslate">SimpleCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell"><code class="docutils literal notranslate">GRUCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.features"><code class="docutils literal notranslate">GRUCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.gate_fn"><code class="docutils literal notranslate">GRUCell.gate_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.activation_fn"><code class="docutils literal notranslate">GRUCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.kernel_init"><code class="docutils literal notranslate">GRUCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.recurrent_kernel_init"><code class="docutils literal notranslate">GRUCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.bias_init"><code class="docutils literal notranslate">GRUCell.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.dtype"><code class="docutils literal notranslate">GRUCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.param_dtype"><code class="docutils literal notranslate">GRUCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.__call__"><code class="docutils literal notranslate">GRUCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.initialize_carry"><code class="docutils literal notranslate">GRUCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell"><code class="docutils literal notranslate">MGUCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.features"><code class="docutils literal notranslate">MGUCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.gate_fn"><code class="docutils literal notranslate">MGUCell.gate_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.activation_fn"><code class="docutils literal notranslate">MGUCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.kernel_init"><code class="docutils literal notranslate">MGUCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.recurrent_kernel_init"><code class="docutils literal notranslate">MGUCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.forget_bias_init"><code class="docutils literal notranslate">MGUCell.forget_bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.activation_bias_init"><code class="docutils literal notranslate">MGUCell.activation_bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.dtype"><code class="docutils literal notranslate">MGUCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.param_dtype"><code class="docutils literal notranslate">MGUCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.reset_gate"><code class="docutils literal notranslate">MGUCell.reset_gate</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.__call__"><code class="docutils literal notranslate">MGUCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.initialize_carry"><code class="docutils literal notranslate">MGUCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN"><code class="docutils literal notranslate">RNN</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.cell"><code class="docutils literal notranslate">RNN.cell</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.time_major"><code class="docutils literal notranslate">RNN.time_major</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.return_carry"><code class="docutils literal notranslate">RNN.return_carry</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.reverse"><code class="docutils literal notranslate">RNN.reverse</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.keep_order"><code class="docutils literal notranslate">RNN.keep_order</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.unroll"><code class="docutils literal notranslate">RNN.unroll</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.variable_axes"><code class="docutils literal notranslate">RNN.variable_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.variable_broadcast"><code class="docutils literal notranslate">RNN.variable_broadcast</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.variable_carry"><code class="docutils literal notranslate">RNN.variable_carry</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.split_rngs"><code class="docutils literal notranslate">RNN.split_rngs</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.__call__"><code class="docutils literal notranslate">RNN.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Bidirectional"><code class="docutils literal notranslate">Bidirectional</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Bidirectional.__call__"><code class="docutils literal notranslate">Bidirectional.__call__()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#batchapply">BatchApply</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.BatchApply"><code class="docutils literal notranslate">BatchApply</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchApply.__call__"><code class="docutils literal notranslate">BatchApply.__call__()</code></a></li> </ul> </li> </ul> </li> </ul> </nav> </div> </div> </div> <div id="searchbox"></div> <article class="bd-article"> <div class="section" id="layers"> <h1>Layers<a class="headerlink" href="#layers" title="Permalink to this heading">#</a></h1> <div class="section" id="linear-modules"> <h2>Linear Modules<a class="headerlink" href="#linear-modules" title="Permalink to this heading">#</a></h2> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.Dense"> class flax.linen.Dense(features, use_bias=True, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, dot_general=None, dot_general_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/linear.html#Dense">[source]</a><a class="headerlink" href="#flax.linen.Dense" title="Permalink to this definition">#</a></dt> <dd>A linear transformation applied over the last dimension of the input. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> layer = nn.Dense(features=4) >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3))) >>> jax.tree_util.tree_map(jnp.shape, params) {'params': {'bias': (4,), 'kernel': (3, 4)}} </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dense.features"> features<a class="headerlink" href="#flax.linen.Dense.features" title="Permalink to this definition">#</a></dt> <dd>the number of output features. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dense.use_bias"> use_bias<a class="headerlink" href="#flax.linen.Dense.use_bias" title="Permalink to this definition">#</a></dt> <dd>whether to add a bias to the output (default: True). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dense.dtype"> dtype<a class="headerlink" href="#flax.linen.Dense.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dense.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.Dense.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dense.precision"> precision<a class="headerlink" href="#flax.linen.Dense.precision" title="Permalink to this definition">#</a></dt> <dd>numerical precision of the computation see <code class="docutils literal notranslate">jax.lax.Precision</code> for details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[None, str, jax._src.lax.lax.Precision, tuple[str, str], tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dense.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.Dense.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the weight matrix. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dense.bias_init"> bias_init<a class="headerlink" href="#flax.linen.Dense.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the bias. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Dense.__call__"> __call__(inputs)<a class="reference internal" href="../../_modules/flax/linen/linear.html#Dense.__call__">[source]</a><a class="headerlink" href="#flax.linen.Dense.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies a linear transformation to the inputs along the last dimension. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd">inputs – The nd-array to be transformed. </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The transformed input. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral"> class flax.linen.DenseGeneral(features, axis=-1, batch_dims=(), use_bias=True, dtype=None, param_dtype=<class 'jax.numpy.float32'>, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, precision=None, dot_general=None, dot_general_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/linear.html#DenseGeneral">[source]</a><a class="headerlink" href="#flax.linen.DenseGeneral" title="Permalink to this definition">#</a></dt> <dd>A linear transformation with flexible axes. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> # equivalent to `nn.Dense(features=4)` >>> layer = nn.DenseGeneral(features=4) >>> # output features (4, 5) >>> layer = nn.DenseGeneral(features=(4, 5)) >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3))) >>> jax.tree_util.tree_map(jnp.shape, params) {'params': {'bias': (4, 5), 'kernel': (3, 4, 5)}} >>> # apply transformation on the the second and last axes >>> layer = nn.DenseGeneral(features=(4, 5), axis=(1, -1)) >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3, 6, 7))) >>> jax.tree_util.tree_map(jnp.shape, params) {'params': {'bias': (4, 5), 'kernel': (3, 7, 4, 5)}} </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.features"> features<a class="headerlink" href="#flax.linen.DenseGeneral.features" title="Permalink to this definition">#</a></dt> <dd>int or tuple with number of output features. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.axis"> axis<a class="headerlink" href="#flax.linen.DenseGeneral.axis" title="Permalink to this definition">#</a></dt> <dd>int or tuple with axes to apply the transformation on. For instance, (-2, -1) will apply the transformation to the last two axes. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.batch_dims"> batch_dims<a class="headerlink" href="#flax.linen.DenseGeneral.batch_dims" title="Permalink to this definition">#</a></dt> <dd>tuple with batch axes. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.use_bias"> use_bias<a class="headerlink" href="#flax.linen.DenseGeneral.use_bias" title="Permalink to this definition">#</a></dt> <dd>whether to add a bias to the output (default: True). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.dtype"> dtype<a class="headerlink" href="#flax.linen.DenseGeneral.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.DenseGeneral.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.DenseGeneral.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the weight matrix. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.bias_init"> bias_init<a class="headerlink" href="#flax.linen.DenseGeneral.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the bias. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.precision"> precision<a class="headerlink" href="#flax.linen.DenseGeneral.precision" title="Permalink to this definition">#</a></dt> <dd>numerical precision of the computation see <code class="docutils literal notranslate">jax.lax.Precision</code> for details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[None, str, jax._src.lax.lax.Precision, tuple[str, str], tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.DenseGeneral.__call__"> __call__(inputs)<a class="reference internal" href="../../_modules/flax/linen/linear.html#DenseGeneral.__call__">[source]</a><a class="headerlink" href="#flax.linen.DenseGeneral.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies a linear transformation to the inputs along multiple dimensions. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd">inputs – The nd-array to be transformed. </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The transformed input. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.Conv"> class flax.linen.Conv(features, kernel_size, strides=1, padding='SAME', input_dilation=1, kernel_dilation=1, feature_group_count=1, use_bias=True, mask=None, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, conv_general_dilated=None, conv_general_dilated_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/linear.html#Conv">[source]</a><a class="headerlink" href="#flax.linen.Conv" title="Permalink to this definition">#</a></dt> <dd>Convolution Module wrapping <code class="docutils literal notranslate">lax.conv_general_dilated</code>. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> # valid padding >>> layer = nn.Conv(features=4, kernel_size=(3,), padding='VALID') >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3))) >>> jax.tree_util.tree_map(jnp.shape, variables) {'params': {'bias': (4,), 'kernel': (3, 3, 4)}} >>> out.shape (1, 6, 4) >>> # circular padding with stride 2 >>> layer = nn.Conv(features=4, kernel_size=(3, 3), strides=2, padding='CIRCULAR') >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3))) >>> jax.tree_util.tree_map(jnp.shape, variables) {'params': {'bias': (4,), 'kernel': (3, 3, 3, 4)}} >>> out.shape (1, 4, 4) >>> # apply lower triangle mask >>> mask = jnp.tril(jnp.ones((3, 3, 4))) >>> layer = nn.Conv(features=4, kernel_size=(3,), mask=mask, padding='VALID') >>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3))) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.features"> features<a class="headerlink" href="#flax.linen.Conv.features" title="Permalink to this definition">#</a></dt> <dd>number of convolution filters. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.kernel_size"> kernel_size<a class="headerlink" href="#flax.linen.Conv.kernel_size" title="Permalink to this definition">#</a></dt> <dd>shape of the convolutional kernel. An integer will be interpreted as a tuple of the single integer. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.strides"> strides<a class="headerlink" href="#flax.linen.Conv.strides" title="Permalink to this definition">#</a></dt> <dd>an integer or a sequence of <cite>n</cite> integers, representing the inter-window strides (default: 1). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">None | int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.padding"> padding<a class="headerlink" href="#flax.linen.Conv.padding" title="Permalink to this definition">#</a></dt> <dd>either the string <code class="docutils literal notranslate">'SAME'</code>, the string <code class="docutils literal notranslate">'VALID'</code>, the string <code class="docutils literal notranslate">'CIRCULAR'</code> (periodic boundary conditions), or a sequence of <code class="docutils literal notranslate">n</code> <code class="docutils literal notranslate">(low, high)</code> integer pairs that give the padding to apply before and after each spatial dimension. A single int is interpreted as applying the same padding in all dims and assign a single int in a sequence causes the same padding to be used on both sides. <code class="docutils literal notranslate">'CAUSAL'</code> padding for a 1D convolution will left-pad the convolution axis, resulting in same-sized output. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, int, collections.abc.Sequence[Union[int, tuple[int, int]]]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.input_dilation"> input_dilation<a class="headerlink" href="#flax.linen.Conv.input_dilation" title="Permalink to this definition">#</a></dt> <dd>an integer or a sequence of <code class="docutils literal notranslate">n</code> integers, giving the dilation factor to apply in each spatial dimension of <code class="docutils literal notranslate">inputs</code> (default: 1). Convolution with input dilation <code class="docutils literal notranslate">d</code> is equivalent to transposed convolution with stride <code class="docutils literal notranslate">d</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">None | int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.kernel_dilation"> kernel_dilation<a class="headerlink" href="#flax.linen.Conv.kernel_dilation" title="Permalink to this definition">#</a></dt> <dd>an integer or a sequence of <code class="docutils literal notranslate">n</code> integers, giving the dilation factor to apply in each spatial dimension of the convolution kernel (default: 1). Convolution with kernel dilation is also known as ‘atrous convolution’. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">None | int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.feature_group_count"> feature_group_count<a class="headerlink" href="#flax.linen.Conv.feature_group_count" title="Permalink to this definition">#</a></dt> <dd>integer, default 1. If specified divides the input features into groups. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.use_bias"> use_bias<a class="headerlink" href="#flax.linen.Conv.use_bias" title="Permalink to this definition">#</a></dt> <dd>whether to add a bias to the output (default: True). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.mask"> mask<a class="headerlink" href="#flax.linen.Conv.mask" title="Permalink to this definition">#</a></dt> <dd>Optional mask for the weights during masked convolution. The mask must be the same shape as the convolution weight matrix. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[jax.Array, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.dtype"> dtype<a class="headerlink" href="#flax.linen.Conv.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.Conv.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.precision"> precision<a class="headerlink" href="#flax.linen.Conv.precision" title="Permalink to this definition">#</a></dt> <dd>numerical precision of the computation see <a href="#id1">``</a>jax.lax.Precision` for details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[None, str, jax._src.lax.lax.Precision, tuple[str, str], tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.Conv.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the convolutional kernel. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Conv.bias_init"> bias_init<a class="headerlink" href="#flax.linen.Conv.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Conv.__call__"> __call__(inputs)<a class="headerlink" href="#flax.linen.Conv.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies a (potentially unshared) convolution to the inputs. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd">inputs – input data with dimensions <code class="docutils literal notranslate">(*batch_dims, spatial_dims..., features)</code>. This is the channels-last convention, i.e. NHWC for a 2d convolution and NDHWC for a 3D convolution. Note: this is different from the input convention used by <code class="docutils literal notranslate">lax.conv_general_dilated</code>, which puts the spatial dimensions last. Note: If the input has more than 1 batch dimension, all batch dimensions are flattened into a single dimension for the convolution and restored before returning. In some cases directly vmap’ing the layer may yield better performance than this default flattening approach. If the input lacks a batch dimension it will be added for the convolution and removed n return, an allowance made to enable writing single-example code. </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The convolved data. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose"> class flax.linen.ConvTranspose(features, kernel_size, strides=None, padding='SAME', kernel_dilation=None, use_bias=True, mask=None, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, transpose_kernel=False, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/linear.html#ConvTranspose">[source]</a><a class="headerlink" href="#flax.linen.ConvTranspose" title="Permalink to this definition">#</a></dt> <dd>Convolution Module wrapping <code class="docutils literal notranslate">lax.conv_transpose</code>. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> # valid padding >>> layer = nn.ConvTranspose(features=4, kernel_size=(3,), padding='VALID') >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3))) >>> jax.tree_util.tree_map(jnp.shape, variables) {'params': {'bias': (4,), 'kernel': (3, 3, 4)}} >>> out.shape (1, 10, 4) >>> # circular padding with stride 2 >>> layer = nn.ConvTranspose(features=4, kernel_size=(6, 6), strides=(2, 2), padding='CIRCULAR', transpose_kernel=True) >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 15, 15, 3))) >>> jax.tree_util.tree_map(jnp.shape, variables) {'params': {'bias': (4,), 'kernel': (6, 6, 4, 3)}} >>> out.shape (1, 30, 30, 4) >>> # apply lower triangle mask >>> mask = jnp.tril(jnp.ones((3, 3, 4))) >>> layer = nn.ConvTranspose(features=4, kernel_size=(3,), mask=mask, padding='VALID') >>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3))) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.features"> features<a class="headerlink" href="#flax.linen.ConvTranspose.features" title="Permalink to this definition">#</a></dt> <dd>number of convolution filters. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.kernel_size"> kernel_size<a class="headerlink" href="#flax.linen.ConvTranspose.kernel_size" title="Permalink to this definition">#</a></dt> <dd>shape of the convolutional kernel. For 1D convolution, the kernel size can be passed as an integer, which will be interpreted as a tuple of the single integer. For all other cases, it must be a sequence of integers. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.strides"> strides<a class="headerlink" href="#flax.linen.ConvTranspose.strides" title="Permalink to this definition">#</a></dt> <dd>an integer or a sequence of <cite>n</cite> integers, representing the inter-window strides. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Sequence[int] | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.padding"> padding<a class="headerlink" href="#flax.linen.ConvTranspose.padding" title="Permalink to this definition">#</a></dt> <dd>either the string <cite>‘SAME’</cite>, the string <cite>‘VALID’</cite>, the string <cite>‘CIRCULAR’</cite> (periodic boundary conditions), or a sequence of <cite>n</cite> <cite>(low, high)</cite> integer pairs that give the padding to apply before and after each spatial dimension. A single int is interpreted as applying the same padding in all dims and assign a single int in a sequence causes the same padding to be used on both sides. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, int, collections.abc.Sequence[Union[int, tuple[int, int]]]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.kernel_dilation"> kernel_dilation<a class="headerlink" href="#flax.linen.ConvTranspose.kernel_dilation" title="Permalink to this definition">#</a></dt> <dd><code class="docutils literal notranslate">None</code>, or an integer or a sequence of <code class="docutils literal notranslate">n</code> integers, giving the dilation factor to apply in each spatial dimension of the convolution kernel. Convolution with kernel dilation is also known as ‘atrous convolution’. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Sequence[int] | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.use_bias"> use_bias<a class="headerlink" href="#flax.linen.ConvTranspose.use_bias" title="Permalink to this definition">#</a></dt> <dd>whether to add a bias to the output (default: True). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.mask"> mask<a class="headerlink" href="#flax.linen.ConvTranspose.mask" title="Permalink to this definition">#</a></dt> <dd>Optional mask for the weights during masked convolution. The mask must be the same shape as the convolution weight matrix. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[jax.Array, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.dtype"> dtype<a class="headerlink" href="#flax.linen.ConvTranspose.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.ConvTranspose.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.precision"> precision<a class="headerlink" href="#flax.linen.ConvTranspose.precision" title="Permalink to this definition">#</a></dt> <dd>numerical precision of the computation see <code class="docutils literal notranslate">jax.lax.Precision</code> for details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[None, str, jax._src.lax.lax.Precision, tuple[str, str], tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.ConvTranspose.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the convolutional kernel. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.bias_init"> bias_init<a class="headerlink" href="#flax.linen.ConvTranspose.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.transpose_kernel"> transpose_kernel<a class="headerlink" href="#flax.linen.ConvTranspose.transpose_kernel" title="Permalink to this definition">#</a></dt> <dd>if <code class="docutils literal notranslate">True</code> flips spatial axes and swaps the input/output channel axes of the kernel. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.ConvTranspose.__call__"> __call__(inputs)<a class="reference internal" href="../../_modules/flax/linen/linear.html#ConvTranspose.__call__">[source]</a><a class="headerlink" href="#flax.linen.ConvTranspose.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies a transposed convolution to the inputs. Behaviour mirrors of <code class="docutils literal notranslate">jax.lax.conv_transpose</code>. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd">inputs – input data with dimensions <code class="docutils literal notranslate">(*batch_dims, spatial_dims..., features).</code> This is the channels-last convention, i.e. NHWC for a 2d convolution and NDHWC for a 3D convolution. Note: this is different from the input convention used by <code class="docutils literal notranslate">lax.conv_general_dilated</code>, which puts the spatial dimensions last. Note: If the input has more than 1 batch dimension, all batch dimensions are flattened into a single dimension for the convolution and restored before returning. In some cases directly vmap’ing the layer may yield better performance than this default flattening approach. If the input lacks a batch dimension it will be added for the convolution and removed n return, an allowance made to enable writing single-example code. </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The convolved data. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.ConvLocal"> class flax.linen.ConvLocal(features, kernel_size, strides=1, padding='SAME', input_dilation=1, kernel_dilation=1, feature_group_count=1, use_bias=True, mask=None, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, conv_general_dilated=None, conv_general_dilated_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/linear.html#ConvLocal">[source]</a><a class="headerlink" href="#flax.linen.ConvLocal" title="Permalink to this definition">#</a></dt> <dd>Local convolution Module wrapping <code class="docutils literal notranslate">lax.conv_general_dilated_local</code>. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> # valid padding >>> layer = nn.ConvLocal(features=4, kernel_size=(3,), padding='VALID') >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3))) >>> jax.tree_util.tree_map(jnp.shape, variables) {'params': {'bias': (6, 4), 'kernel': (6, 9, 4)}} >>> out.shape (1, 6, 4) >>> # circular padding with stride 2 >>> layer = nn.ConvLocal(features=4, kernel_size=(3, 3), strides=2, padding='CIRCULAR') >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3))) >>> jax.tree_util.tree_map(jnp.shape, variables) {'params': {'bias': (1, 4, 4), 'kernel': (1, 4, 27, 4)}} >>> out.shape (1, 4, 4) >>> # apply lower triangle mask >>> mask = jnp.tril(jnp.ones((6, 9, 4))) >>> layer = nn.ConvLocal(features=4, kernel_size=(3,), mask=mask, padding='VALID') >>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3))) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.features"> features<a class="headerlink" href="#flax.linen.ConvLocal.features" title="Permalink to this definition">#</a></dt> <dd>number of convolution filters. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.kernel_size"> kernel_size<a class="headerlink" href="#flax.linen.ConvLocal.kernel_size" title="Permalink to this definition">#</a></dt> <dd>shape of the convolutional kernel. An integer will be interpreted as a tuple of the single integer. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.strides"> strides<a class="headerlink" href="#flax.linen.ConvLocal.strides" title="Permalink to this definition">#</a></dt> <dd>an integer or a sequence of <cite>n</cite> integers, representing the inter-window strides (default: 1). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">None | int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.padding"> padding<a class="headerlink" href="#flax.linen.ConvLocal.padding" title="Permalink to this definition">#</a></dt> <dd>either the string <code class="docutils literal notranslate">'SAME'</code>, the string <code class="docutils literal notranslate">'VALID'</code>, the string <code class="docutils literal notranslate">'CIRCULAR'</code> (periodic boundary conditions), or a sequence of <code class="docutils literal notranslate">n</code> <code class="docutils literal notranslate">(low, high)</code> integer pairs that give the padding to apply before and after each spatial dimension. A single int is interpreted as applying the same padding in all dims and assign a single int in a sequence causes the same padding to be used on both sides. <code class="docutils literal notranslate">'CAUSAL'</code> padding for a 1D convolution will left-pad the convolution axis, resulting in same-sized output. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, int, collections.abc.Sequence[Union[int, tuple[int, int]]]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.input_dilation"> input_dilation<a class="headerlink" href="#flax.linen.ConvLocal.input_dilation" title="Permalink to this definition">#</a></dt> <dd>an integer or a sequence of <code class="docutils literal notranslate">n</code> integers, giving the dilation factor to apply in each spatial dimension of <code class="docutils literal notranslate">inputs</code> (default: 1). Convolution with input dilation <code class="docutils literal notranslate">d</code> is equivalent to transposed convolution with stride <code class="docutils literal notranslate">d</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">None | int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.kernel_dilation"> kernel_dilation<a class="headerlink" href="#flax.linen.ConvLocal.kernel_dilation" title="Permalink to this definition">#</a></dt> <dd>an integer or a sequence of <code class="docutils literal notranslate">n</code> integers, giving the dilation factor to apply in each spatial dimension of the convolution kernel (default: 1). Convolution with kernel dilation is also known as ‘atrous convolution’. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">None | int | collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.feature_group_count"> feature_group_count<a class="headerlink" href="#flax.linen.ConvLocal.feature_group_count" title="Permalink to this definition">#</a></dt> <dd>integer, default 1. If specified divides the input features into groups. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.use_bias"> use_bias<a class="headerlink" href="#flax.linen.ConvLocal.use_bias" title="Permalink to this definition">#</a></dt> <dd>whether to add a bias to the output (default: True). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.mask"> mask<a class="headerlink" href="#flax.linen.ConvLocal.mask" title="Permalink to this definition">#</a></dt> <dd>Optional mask for the weights during masked convolution. The mask must be the same shape as the convolution weight matrix. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[jax.Array, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.dtype"> dtype<a class="headerlink" href="#flax.linen.ConvLocal.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.ConvLocal.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.precision"> precision<a class="headerlink" href="#flax.linen.ConvLocal.precision" title="Permalink to this definition">#</a></dt> <dd>numerical precision of the computation see <code class="docutils literal notranslate">jax.lax.Precision</code> for details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[None, str, jax._src.lax.lax.Precision, tuple[str, str], tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.ConvLocal.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the convolutional kernel. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.bias_init"> bias_init<a class="headerlink" href="#flax.linen.ConvLocal.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.ConvLocal.__call__"> __call__(inputs)<a class="headerlink" href="#flax.linen.ConvLocal.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies a (potentially unshared) convolution to the inputs. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd">inputs – input data with dimensions <code class="docutils literal notranslate">(*batch_dims, spatial_dims..., features)</code>. This is the channels-last convention, i.e. NHWC for a 2d convolution and NDHWC for a 3D convolution. Note: this is different from the input convention used by <code class="docutils literal notranslate">lax.conv_general_dilated</code>, which puts the spatial dimensions last. Note: If the input has more than 1 batch dimension, all batch dimensions are flattened into a single dimension for the convolution and restored before returning. In some cases directly vmap’ing the layer may yield better performance than this default flattening approach. If the input lacks a batch dimension it will be added for the convolution and removed n return, an allowance made to enable writing single-example code. </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The convolved data. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.Einsum"> class flax.linen.Einsum(shape, einsum_str=None, use_bias=True, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/linear.html#Einsum">[source]</a><a class="headerlink" href="#flax.linen.Einsum" title="Permalink to this definition">#</a></dt> <dd>An einsum transformation with learnable kernel and bias. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> layer = nn.Einsum((5, 6, 7), 'abc,cde->abde') >>> variables = layer.init(jax.random.key(0), jnp.ones((3, 4, 5))) >>> jax.tree_util.tree_map(jnp.shape, variables) {'params': {'bias': (6, 7), 'kernel': (5, 6, 7)}} </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Einsum.shape"> shape<a class="headerlink" href="#flax.linen.Einsum.shape" title="Permalink to this definition">#</a></dt> <dd>the shape of the kernel. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Einsum.einsum_str"> einsum_str<a class="headerlink" href="#flax.linen.Einsum.einsum_str" title="Permalink to this definition">#</a></dt> <dd>a string to denote the einsum equation. The equation must have exactly two operands, the lhs being the input passed in, and the rhs being the learnable kernel. Exactly one of <code class="docutils literal notranslate">einsum_str</code> in the constructor argument and call argument must be not None, while the other must be None. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">str | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Einsum.use_bias"> use_bias<a class="headerlink" href="#flax.linen.Einsum.use_bias" title="Permalink to this definition">#</a></dt> <dd>whether to add a bias to the output (default: True). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Einsum.dtype"> dtype<a class="headerlink" href="#flax.linen.Einsum.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Einsum.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.Einsum.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Einsum.precision"> precision<a class="headerlink" href="#flax.linen.Einsum.precision" title="Permalink to this definition">#</a></dt> <dd>numerical precision of the computation see <code class="docutils literal notranslate">jax.lax.Precision</code> for details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[None, str, jax._src.lax.lax.Precision, tuple[str, str], tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Einsum.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.Einsum.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the weight matrix. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Einsum.bias_init"> bias_init<a class="headerlink" href="#flax.linen.Einsum.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the bias. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Einsum.__call__"> __call__(inputs, einsum_str=None)<a class="reference internal" href="../../_modules/flax/linen/linear.html#Einsum.__call__">[source]</a><a class="headerlink" href="#flax.linen.Einsum.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies a linear transformation to the inputs along the last dimension. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>inputs – The nd-array to be transformed.</li> <li>einsum_str – a string to denote the einsum equation. The equation must have exactly two operands, the lhs being the input passed in, and the rhs being the learnable kernel. The <code class="docutils literal notranslate">einsum_str</code> passed into the call method will take precedence over the <code class="docutils literal notranslate">einsum_str</code> passed into the constructor.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The transformed input. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.Embed"> class flax.linen.Embed(num_embeddings, features, dtype=None, param_dtype=<class 'jax.numpy.float32'>, embedding_init=<function variance_scaling.<locals>.init>, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/linear.html#Embed">[source]</a><a class="headerlink" href="#flax.linen.Embed" title="Permalink to this definition">#</a></dt> <dd>Embedding Module. A parameterized function from integers [0, <code class="docutils literal notranslate">num_embeddings</code>) to <code class="docutils literal notranslate">features</code>-dimensional vectors. This <code class="docutils literal notranslate">Module</code> will create an <code class="docutils literal notranslate">embedding</code> matrix with shape <code class="docutils literal notranslate">(num_embeddings, features)</code>. When calling this layer, the input values will be used to 0-index into the <code class="docutils literal notranslate">embedding</code> matrix. Indexing on a value greater than or equal to <code class="docutils literal notranslate">num_embeddings</code> will result in <code class="docutils literal notranslate">nan</code> values. When <code class="docutils literal notranslate">num_embeddings</code> equals to 1, it will broadcast the <code class="docutils literal notranslate">embedding</code> matrix to input shape with <code class="docutils literal notranslate">features</code> dimension appended. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> layer = nn.Embed(num_embeddings=5, features=3) >>> indices_input = jnp.array([[0, 1, 2], [-1, -2, -3]]) >>> variables = layer.init(jax.random.key(0), indices_input) >>> variables {'params': {'embedding': Array([[-0.28884724, 0.19018005, -0.414205 ], [-0.11768015, -0.54618824, -0.3789283 ], [ 0.30428642, 0.49511626, 0.01706631], [-0.0982546 , -0.43055868, 0.20654906], [-0.688412 , -0.46882293, 0.26723292]], dtype=float32)}} >>> # get the first three and last three embeddings >>> layer.apply(variables, indices_input) Array([[[-0.28884724, 0.19018005, -0.414205 ], [-0.11768015, -0.54618824, -0.3789283 ], [ 0.30428642, 0.49511626, 0.01706631]], [[-0.688412 , -0.46882293, 0.26723292], [-0.0982546 , -0.43055868, 0.20654906], [ 0.30428642, 0.49511626, 0.01706631]]], dtype=float32) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Embed.num_embeddings"> num_embeddings<a class="headerlink" href="#flax.linen.Embed.num_embeddings" title="Permalink to this definition">#</a></dt> <dd>number of embeddings / vocab size. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Embed.features"> features<a class="headerlink" href="#flax.linen.Embed.features" title="Permalink to this definition">#</a></dt> <dd>number of feature dimensions for each embedding. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Embed.dtype"> dtype<a class="headerlink" href="#flax.linen.Embed.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the embedding vectors (default: same as embedding). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Embed.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.Embed.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Embed.embedding_init"> embedding_init<a class="headerlink" href="#flax.linen.Embed.embedding_init" title="Permalink to this definition">#</a></dt> <dd>embedding initializer. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Embed.__call__"> __call__(inputs)<a class="reference internal" href="../../_modules/flax/linen/linear.html#Embed.__call__">[source]</a><a class="headerlink" href="#flax.linen.Embed.__call__" title="Permalink to this definition">#</a></dt> <dd>Embeds the inputs along the last dimension. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd">inputs – input data, all dimensions are considered batch dimensions. Values in the input array must be integers. </dd> <dt class="field-even">Returns</dt> <dd class="field-even">Output which is embedded input data. The output shape follows the input, with an additional <code class="docutils literal notranslate">features</code> dimension appended. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Embed.attend"> attend(query)<a class="reference internal" href="../../_modules/flax/linen/linear.html#Embed.attend">[source]</a><a class="headerlink" href="#flax.linen.Embed.attend" title="Permalink to this definition">#</a></dt> <dd>Attend over the embedding using a query array. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd">query – array with last dimension equal the feature depth <code class="docutils literal notranslate">features</code> of the embedding. </dd> <dt class="field-even">Returns</dt> <dd class="field-even">An array with final dim <code class="docutils literal notranslate">num_embeddings</code> corresponding to the batched inner-product of the array of query vectors against each embedding. Commonly used for weight-sharing between embeddings and logit transform in NLP models. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> <tr class="row-odd"><td><a class="reference internal" href="#flax.linen.Embed.attend" title="flax.linen.Embed.attend"><code class="xref py py-obj docutils literal notranslate">attend</code></a>(query)</td> <td>Attend over the embedding using a query array.</td> </tr> </tbody> </table> </div> </dd></dl> </div> </div> <div class="section" id="pooling"> <h2>Pooling<a class="headerlink" href="#pooling" title="Permalink to this heading">#</a></h2> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.max_pool"> flax.linen.max_pool(inputs, window_shape, strides=None, padding='VALID')<a class="reference internal" href="../../_modules/flax/linen/pooling.html#max_pool">[source]</a><a class="headerlink" href="#flax.linen.max_pool" title="Permalink to this definition">#</a></dt> <dd>Pools the input by taking the maximum of a window slice. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>inputs – input data with dimensions (batch, window dims…, features).</li> <li>window_shape – a shape tuple defining the window to reduce over.</li> <li>strides – a sequence of <code class="docutils literal notranslate">n</code> integers, representing the inter-window strides (default: <code class="docutils literal notranslate">(1, ..., 1)</code>).</li> <li>padding – either the string <code class="docutils literal notranslate">'SAME'</code>, the string <code class="docutils literal notranslate">'VALID'</code>, or a sequence of <code class="docutils literal notranslate">n</code> <code class="docutils literal notranslate">(low, high)</code> integer pairs that give the padding to apply before and after each spatial dimension (default: <code class="docutils literal notranslate">'VALID'</code>).</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The maximum for each window slice. </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.avg_pool"> flax.linen.avg_pool(inputs, window_shape, strides=None, padding='VALID', count_include_pad=True)<a class="reference internal" href="../../_modules/flax/linen/pooling.html#avg_pool">[source]</a><a class="headerlink" href="#flax.linen.avg_pool" title="Permalink to this definition">#</a></dt> <dd>Pools the input by taking the average over a window. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>inputs – input data with dimensions (batch, window dims…, features).</li> <li>window_shape – a shape tuple defining the window to reduce over.</li> <li>strides – a sequence of <code class="docutils literal notranslate">n</code> integers, representing the inter-window strides (default: <code class="docutils literal notranslate">(1, ..., 1)</code>).</li> <li>padding – either the string <code class="docutils literal notranslate">'SAME'</code>, the string <code class="docutils literal notranslate">'VALID'</code>, or a sequence of <code class="docutils literal notranslate">n</code> <code class="docutils literal notranslate">(low, high)</code> integer pairs that give the padding to apply before and after each spatial dimension (default: <code class="docutils literal notranslate">'VALID'</code>).</li> <li>count_include_pad – a boolean whether to include padded tokens in the average calculation (default: <code class="docutils literal notranslate">True</code>).</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The average for each window slice. </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.pool"> flax.linen.pool(inputs, init, reduce_fn, window_shape, strides, padding)<a class="reference internal" href="../../_modules/flax/linen/pooling.html#pool">[source]</a><a class="headerlink" href="#flax.linen.pool" title="Permalink to this definition">#</a></dt> <dd>Helper function to define pooling functions. Pooling functions are implemented using the ReduceWindow XLA op. <div class="admonition note"> Note Be aware that pooling is not generally differentiable. That means providing a reduce_fn that is differentiable does not imply that pool is differentiable. </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>inputs – input data with dimensions (batch, window dims…, features).</li> <li>init – the initial value for the reduction</li> <li>reduce_fn – a reduce function of the form <code class="docutils literal notranslate">(T, T) -> T</code>.</li> <li>window_shape – a shape tuple defining the window to reduce over.</li> <li>strides – a sequence of <code class="docutils literal notranslate">n</code> integers, representing the inter-window strides (default: <code class="docutils literal notranslate">(1, ..., 1)</code>).</li> <li>padding – either the string <code class="docutils literal notranslate">'SAME'</code>, the string <code class="docutils literal notranslate">'VALID'</code>, or a sequence of <code class="docutils literal notranslate">n</code> <code class="docutils literal notranslate">(low, high)</code> integer pairs that give the padding to apply before and after each spatial dimension.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The output of the reduction for each window slice. </dd> </dl> </dd></dl> </div> <div class="section" id="normalization"> <h2>Normalization<a class="headerlink" href="#normalization" title="Permalink to this heading">#</a></h2> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.BatchNorm"> class flax.linen.BatchNorm(use_running_average=None, axis=-1, momentum=0.99, epsilon=1e-05, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_bias=True, use_scale=True, bias_init=<function zeros>, scale_init=<function ones>, axis_name=None, axis_index_groups=None, use_fast_variance=True, force_float32_reductions=True, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#BatchNorm">[source]</a><a class="headerlink" href="#flax.linen.BatchNorm" title="Permalink to this definition">#</a></dt> <dd>BatchNorm Module. Usage Note: If we define a model with BatchNorm, for example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> BN = nn.BatchNorm(momentum=0.9, epsilon=1e-5, dtype=jnp.float32) </pre></div> </div> The initialized variables dict will contain, in addition to a ‘params’ collection, a separate ‘batch_stats’ collection that will contain all the running statistics for all the BatchNorm layers in a model: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> x = jax.random.normal(jax.random.key(0), (5, 6)) >>> variables = BN.init(jax.random.key(1), x, use_running_average=False) >>> jax.tree_util.tree_map(jnp.shape, variables) {'batch_stats': {'mean': (6,), 'var': (6,)}, 'params': {'bias': (6,), 'scale': (6,)}} </pre></div> </div> We then update the batch_stats during training by specifying that the <code class="docutils literal notranslate">batch_stats</code> collection is mutable in the <code class="docutils literal notranslate">apply</code> method for our module.: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> y, new_batch_stats = BN.apply(variables, x, mutable=['batch_stats'], use_running_average=False) </pre></div> </div> During eval we would define BN with <code class="docutils literal notranslate">use_running_average=True</code> and use the batch_stats collection from training to set the statistics. In this case we are not mutating the batch statistics collection, and needn’t mark it mutable: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> y = BN.apply(variables, x, mutable=['batch_stats'], use_running_average=True) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.use_running_average"> use_running_average<a class="headerlink" href="#flax.linen.BatchNorm.use_running_average" title="Permalink to this definition">#</a></dt> <dd>if True, the statistics stored in batch_stats will be used instead of computing the batch statistics on the input. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.axis"> axis<a class="headerlink" href="#flax.linen.BatchNorm.axis" title="Permalink to this definition">#</a></dt> <dd>the feature or non-batch axis of the input. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.momentum"> momentum<a class="headerlink" href="#flax.linen.BatchNorm.momentum" title="Permalink to this definition">#</a></dt> <dd>decay rate for the exponential moving average of the batch statistics. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.epsilon"> epsilon<a class="headerlink" href="#flax.linen.BatchNorm.epsilon" title="Permalink to this definition">#</a></dt> <dd>a small float added to variance to avoid dividing by zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.dtype"> dtype<a class="headerlink" href="#flax.linen.BatchNorm.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the result (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.BatchNorm.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.use_bias"> use_bias<a class="headerlink" href="#flax.linen.BatchNorm.use_bias" title="Permalink to this definition">#</a></dt> <dd>if True, bias (beta) is added. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.use_scale"> use_scale<a class="headerlink" href="#flax.linen.BatchNorm.use_scale" title="Permalink to this definition">#</a></dt> <dd>if True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.bias_init"> bias_init<a class="headerlink" href="#flax.linen.BatchNorm.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for bias, by default, zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.scale_init"> scale_init<a class="headerlink" href="#flax.linen.BatchNorm.scale_init" title="Permalink to this definition">#</a></dt> <dd>initializer for scale, by default, one. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.axis_name"> axis_name<a class="headerlink" href="#flax.linen.BatchNorm.axis_name" title="Permalink to this definition">#</a></dt> <dd>the axis name used to combine batch statistics from multiple devices. See <code class="docutils literal notranslate">jax.pmap</code> for a description of axis names (default: None). Note, this is only used for pmap and shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">str | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.axis_index_groups"> axis_index_groups<a class="headerlink" href="#flax.linen.BatchNorm.axis_index_groups" title="Permalink to this definition">#</a></dt> <dd>groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, <code class="docutils literal notranslate">[[0, 1], [2, 3]]</code> would independently batch-normalize over the examples on the first two and last two devices. See <code class="docutils literal notranslate">jax.lax.psum</code> for more details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Any </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.use_fast_variance"> use_fast_variance<a class="headerlink" href="#flax.linen.BatchNorm.use_fast_variance" title="Permalink to this definition">#</a></dt> <dd>If true, use a faster, but less numerically stable, calculation for the variance. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.BatchNorm.__call__"> __call__(x, use_running_average=None, *, mask=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#BatchNorm.__call__">[source]</a><a class="headerlink" href="#flax.linen.BatchNorm.__call__" title="Permalink to this definition">#</a></dt> <dd>Normalizes the input using batch statistics. <div class="admonition note"> Note During initialization (when <code class="docutils literal notranslate">self.is_initializing()</code> is <code class="docutils literal notranslate">True</code>) the running average of the batch statistics will not be updated. Therefore, the inputs fed during initialization don’t need to match that of the actual input distribution and the reduction axis (set with <code class="docutils literal notranslate">axis_name</code>) does not have to exist. </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>x – the input to be normalized.</li> <li>use_running_average – if true, the statistics stored in batch_stats will be used instead of computing the batch statistics on the input.</li> <li>mask – Binary array of shape broadcastable to <code class="docutils literal notranslate">inputs</code> tensor, indicating the positions for which the mean and variance should be computed.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">Normalized inputs (the same shape as inputs). </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.LayerNorm"> class flax.linen.LayerNorm(epsilon=1e-06, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_bias=True, use_scale=True, bias_init=<function zeros>, scale_init=<function ones>, reduction_axes=-1, feature_axes=-1, axis_name=None, axis_index_groups=None, use_fast_variance=True, force_float32_reductions=True, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#LayerNorm">[source]</a><a class="headerlink" href="#flax.linen.LayerNorm" title="Permalink to this definition">#</a></dt> <dd>Layer normalization (<a class="reference external" href="https://arxiv.org/abs/1607.06450">https://arxiv.org/abs/1607.06450</a>). LayerNorm normalizes the activations of the layer for each given example in a batch independently, rather than across a batch like Batch Normalization. i.e. applies a transformation that maintains the mean activation within each example close to 0 and the activation standard deviation close to 1. <div class="admonition note"> Note This normalization operation is identical to InstanceNorm and GroupNorm; the difference is simply which axes are reduced and the shape of the feature axes (i.e. the shape of the learnable scale and bias parameters). </div> Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax >>> import numpy as np >>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6)) >>> layer = nn.LayerNorm() >>> variables = layer.init(jax.random.key(1), x) >>> variables {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}} >>> y = layer.apply(variables, x) >>> y = nn.LayerNorm(reduction_axes=(1, 2, 3)).apply(variables, x) >>> y2 = nn.GroupNorm(num_groups=1).apply(variables, x) >>> np.testing.assert_allclose(y, y2) >>> y = nn.LayerNorm(reduction_axes=(1, 2), feature_axes=-1).apply(variables, x) >>> y2 = nn.InstanceNorm(feature_axes=-1).apply(variables, x) >>> np.testing.assert_allclose(y, y2) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.epsilon"> epsilon<a class="headerlink" href="#flax.linen.LayerNorm.epsilon" title="Permalink to this definition">#</a></dt> <dd>A small float added to variance to avoid dividing by zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.dtype"> dtype<a class="headerlink" href="#flax.linen.LayerNorm.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the result (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.LayerNorm.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.use_bias"> use_bias<a class="headerlink" href="#flax.linen.LayerNorm.use_bias" title="Permalink to this definition">#</a></dt> <dd>If True, bias (beta) is added. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.use_scale"> use_scale<a class="headerlink" href="#flax.linen.LayerNorm.use_scale" title="Permalink to this definition">#</a></dt> <dd>If True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.bias_init"> bias_init<a class="headerlink" href="#flax.linen.LayerNorm.bias_init" title="Permalink to this definition">#</a></dt> <dd>Initializer for bias, by default, zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.scale_init"> scale_init<a class="headerlink" href="#flax.linen.LayerNorm.scale_init" title="Permalink to this definition">#</a></dt> <dd>Initializer for scale, by default, one. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.reduction_axes"> reduction_axes<a class="headerlink" href="#flax.linen.LayerNorm.reduction_axes" title="Permalink to this definition">#</a></dt> <dd>Axes for computing normalization statistics. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[int, collections.abc.Sequence[int]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.feature_axes"> feature_axes<a class="headerlink" href="#flax.linen.LayerNorm.feature_axes" title="Permalink to this definition">#</a></dt> <dd>Feature axes for learned bias and scaling. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[int, collections.abc.Sequence[int]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.axis_name"> axis_name<a class="headerlink" href="#flax.linen.LayerNorm.axis_name" title="Permalink to this definition">#</a></dt> <dd>the axis name used to combine batch statistics from multiple devices. See <code class="docutils literal notranslate">jax.pmap</code> for a description of axis names (default: None). This is only needed if the model is subdivided across devices, i.e. the array being normalized is sharded across devices within a pmap or shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">str | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.axis_index_groups"> axis_index_groups<a class="headerlink" href="#flax.linen.LayerNorm.axis_index_groups" title="Permalink to this definition">#</a></dt> <dd>groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, <code class="docutils literal notranslate">[[0, 1], [2, 3]]</code> would independently batch-normalize over the examples on the first two and last two devices. See <code class="docutils literal notranslate">jax.lax.psum</code> for more details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Any </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.use_fast_variance"> use_fast_variance<a class="headerlink" href="#flax.linen.LayerNorm.use_fast_variance" title="Permalink to this definition">#</a></dt> <dd>If true, use a faster, but less numerically stable, calculation for the variance. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.LayerNorm.__call__"> __call__(x, *, mask=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#LayerNorm.__call__">[source]</a><a class="headerlink" href="#flax.linen.LayerNorm.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies layer normalization on the input. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>x – the inputs</li> <li>mask – Binary array of shape broadcastable to <code class="docutils literal notranslate">inputs</code> tensor, indicating the positions for which the mean and variance should be computed.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">Normalized inputs (the same shape as inputs). </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.GroupNorm"> class flax.linen.GroupNorm(num_groups=32, group_size=None, epsilon=1e-06, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_bias=True, use_scale=True, bias_init=<function zeros>, scale_init=<function ones>, reduction_axes=None, axis_name=None, axis_index_groups=None, use_fast_variance=True, force_float32_reductions=True, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#GroupNorm">[source]</a><a class="headerlink" href="#flax.linen.GroupNorm" title="Permalink to this definition">#</a></dt> <dd>Group normalization (arxiv.org/abs/1803.08494). This op is similar to batch normalization, but statistics are shared across equally-sized groups of channels and not shared across batch dimension. Thus, group normalization does not depend on the batch composition and does not require maintaining internal state for storing statistics. The user should either specify the total number of channel groups or the number of channels per group. <div class="admonition note"> Note LayerNorm is a special case of GroupNorm where <code class="docutils literal notranslate">num_groups=1</code>, and InstanceNorm is a special case of GroupNorm where <code class="docutils literal notranslate">group_size=1</code>. </div> Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax >>> import numpy as np >>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6)) >>> layer = nn.GroupNorm(num_groups=3) >>> variables = layer.init(jax.random.key(1), x) >>> variables {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}} >>> y = layer.apply(variables, x) >>> y = nn.GroupNorm(num_groups=1).apply(variables, x) >>> y2 = nn.LayerNorm(reduction_axes=(1, 2, 3)).apply(variables, x) >>> np.testing.assert_allclose(y, y2) >>> y = nn.GroupNorm(num_groups=None, group_size=1).apply(variables, x) >>> y2 = nn.InstanceNorm(feature_axes=-1).apply(variables, x) >>> np.testing.assert_allclose(y, y2) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.num_groups"> num_groups<a class="headerlink" href="#flax.linen.GroupNorm.num_groups" title="Permalink to this definition">#</a></dt> <dd>the total number of channel groups. The default value of 32 is proposed by the original group normalization paper. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.group_size"> group_size<a class="headerlink" href="#flax.linen.GroupNorm.group_size" title="Permalink to this definition">#</a></dt> <dd>the number of channels in a group. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.epsilon"> epsilon<a class="headerlink" href="#flax.linen.GroupNorm.epsilon" title="Permalink to this definition">#</a></dt> <dd>A small float added to variance to avoid dividing by zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.dtype"> dtype<a class="headerlink" href="#flax.linen.GroupNorm.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the result (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.GroupNorm.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.use_bias"> use_bias<a class="headerlink" href="#flax.linen.GroupNorm.use_bias" title="Permalink to this definition">#</a></dt> <dd>If True, bias (beta) is added. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.use_scale"> use_scale<a class="headerlink" href="#flax.linen.GroupNorm.use_scale" title="Permalink to this definition">#</a></dt> <dd>If True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.bias_init"> bias_init<a class="headerlink" href="#flax.linen.GroupNorm.bias_init" title="Permalink to this definition">#</a></dt> <dd>Initializer for bias, by default, zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.scale_init"> scale_init<a class="headerlink" href="#flax.linen.GroupNorm.scale_init" title="Permalink to this definition">#</a></dt> <dd>Initializer for scale, by default, one. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.reduction_axes"> reduction_axes<a class="headerlink" href="#flax.linen.GroupNorm.reduction_axes" title="Permalink to this definition">#</a></dt> <dd>List of axes used for computing normalization statistics. This list must include the final dimension, which is assumed to be the feature axis. Furthermore, if the input used at call time has additional leading axes compared to the data used for initialisation, for example due to batching, then the reduction axes need to be defined explicitly. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[int, collections.abc.Sequence[int]]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.axis_name"> axis_name<a class="headerlink" href="#flax.linen.GroupNorm.axis_name" title="Permalink to this definition">#</a></dt> <dd>the axis name used to combine batch statistics from multiple devices. See <code class="docutils literal notranslate">jax.pmap</code> for a description of axis names (default: None). This is only needed if the model is subdivided across devices, i.e. the array being normalized is sharded across devices within a pmap or shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">str | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.axis_index_groups"> axis_index_groups<a class="headerlink" href="#flax.linen.GroupNorm.axis_index_groups" title="Permalink to this definition">#</a></dt> <dd>groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, <code class="docutils literal notranslate">[[0, 1], [2, 3]]</code> would independently batch-normalize over the examples on the first two and last two devices. See <code class="docutils literal notranslate">jax.lax.psum</code> for more details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Any </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.use_fast_variance"> use_fast_variance<a class="headerlink" href="#flax.linen.GroupNorm.use_fast_variance" title="Permalink to this definition">#</a></dt> <dd>If true, use a faster, but less numerically stable, calculation for the variance. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.GroupNorm.__call__"> __call__(x, *, mask=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#GroupNorm.__call__">[source]</a><a class="headerlink" href="#flax.linen.GroupNorm.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies group normalization to the input (arxiv.org/abs/1803.08494). <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>x – the input of shape <code class="docutils literal notranslate">...C</code> where <code class="docutils literal notranslate">C</code> is a channels dimension and <code class="docutils literal notranslate">...</code> represents an arbitrary number of extra dimensions that can be used to accumulate statistics over. If no reduction axes have been specified then all additional dimensions <code class="docutils literal notranslate">...</code> will be used to accumulate statistics apart from the leading dimension which is assumed to represent the batch.</li> <li>mask – Binary array of shape broadcastable to <code class="docutils literal notranslate">inputs</code> tensor, indicating the positions for which the mean and variance should be computed.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">Normalized inputs (the same shape as inputs). </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.RMSNorm"> class flax.linen.RMSNorm(epsilon=1e-06, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_scale=True, scale_init=<function ones>, reduction_axes=-1, feature_axes=-1, axis_name=None, axis_index_groups=None, use_fast_variance=True, force_float32_reductions=True, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#RMSNorm">[source]</a><a class="headerlink" href="#flax.linen.RMSNorm" title="Permalink to this definition">#</a></dt> <dd>RMS Layer normalization (<a class="reference external" href="https://arxiv.org/abs/1910.07467">https://arxiv.org/abs/1910.07467</a>). RMSNorm normalizes the activations of the layer for each given example in a batch independently, rather than across a batch like Batch Normalization. Unlike LayerNorm which re-centers the mean to be 0 and normalizes by the standard deviation of the activations, RMSNorm does not re-center at all and instead normalizes by the root mean square of the activations. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax >>> x = jax.random.normal(jax.random.key(0), (5, 6)) >>> layer = nn.RMSNorm() >>> variables = layer.init(jax.random.key(1), x) >>> variables {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32)}} >>> y = layer.apply(variables, x) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.epsilon"> epsilon<a class="headerlink" href="#flax.linen.RMSNorm.epsilon" title="Permalink to this definition">#</a></dt> <dd>A small float added to variance to avoid dividing by zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.dtype"> dtype<a class="headerlink" href="#flax.linen.RMSNorm.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the result (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.RMSNorm.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.use_scale"> use_scale<a class="headerlink" href="#flax.linen.RMSNorm.use_scale" title="Permalink to this definition">#</a></dt> <dd>If True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.scale_init"> scale_init<a class="headerlink" href="#flax.linen.RMSNorm.scale_init" title="Permalink to this definition">#</a></dt> <dd>Initializer for scale, by default, one. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.reduction_axes"> reduction_axes<a class="headerlink" href="#flax.linen.RMSNorm.reduction_axes" title="Permalink to this definition">#</a></dt> <dd>Axes for computing normalization statistics. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[int, collections.abc.Sequence[int]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.feature_axes"> feature_axes<a class="headerlink" href="#flax.linen.RMSNorm.feature_axes" title="Permalink to this definition">#</a></dt> <dd>Feature axes for learned bias and scaling. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[int, collections.abc.Sequence[int]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.axis_name"> axis_name<a class="headerlink" href="#flax.linen.RMSNorm.axis_name" title="Permalink to this definition">#</a></dt> <dd>the axis name used to combine batch statistics from multiple devices. See <code class="docutils literal notranslate">jax.pmap</code> for a description of axis names (default: None). This is only needed if the model is subdivided across devices, i.e. the array being normalized is sharded across devices within a pmap or shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">str | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.axis_index_groups"> axis_index_groups<a class="headerlink" href="#flax.linen.RMSNorm.axis_index_groups" title="Permalink to this definition">#</a></dt> <dd>groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, <code class="docutils literal notranslate">[[0, 1], [2, 3]]</code> would independently batch-normalize over the examples on the first two and last two devices. See <code class="docutils literal notranslate">jax.lax.psum</code> for more details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Any </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.use_fast_variance"> use_fast_variance<a class="headerlink" href="#flax.linen.RMSNorm.use_fast_variance" title="Permalink to this definition">#</a></dt> <dd>If true, use a faster, but less numerically stable, calculation for the variance. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.RMSNorm.__call__"> __call__(x, *, mask=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#RMSNorm.__call__">[source]</a><a class="headerlink" href="#flax.linen.RMSNorm.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies RMS layer normalization on the input. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>x – the inputs</li> <li>mask – Binary array of shape broadcastable to <code class="docutils literal notranslate">inputs</code> tensor, indicating the positions for which the mean and variance should be computed.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">Normalized inputs (the same shape as inputs). </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm"> class flax.linen.InstanceNorm(epsilon=1e-06, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_bias=True, use_scale=True, bias_init=<function zeros>, scale_init=<function ones>, feature_axes=-1, axis_name=None, axis_index_groups=None, use_fast_variance=True, force_float32_reductions=True, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#InstanceNorm">[source]</a><a class="headerlink" href="#flax.linen.InstanceNorm" title="Permalink to this definition">#</a></dt> <dd>Instance normalization (<a class="reference external" href="https://arxiv.org/abs/1607.08022v3">https://arxiv.org/abs/1607.08022v3</a>). InstanceNorm normalizes the activations of the layer for each channel (rather than across all channels like Layer Normalization), and for each given example in a batch independently (rather than across an entire batch like Batch Normalization). i.e. applies a transformation that maintains the mean activation within each channel within each example close to 0 and the activation standard deviation close to 1. <div class="admonition note"> Note This normalization operation is identical to LayerNorm and GroupNorm; the difference is simply which axes are reduced and the shape of the feature axes (i.e. the shape of the learnable scale and bias parameters). </div> Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax >>> import numpy as np >>> # dimensions: (batch, height, width, channel) >>> x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5)) >>> layer = nn.InstanceNorm() >>> variables = layer.init(jax.random.key(1), x) >>> variables {'params': {'scale': Array([1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0.], dtype=float32)}} >>> y = layer.apply(variables, x) >>> # having a channel_axis of -1 in InstanceNorm is identical to reducing all non-batch, >>> # non-channel axes and using the feature_axes as the feature_axes in LayerNorm >>> y2 = nn.LayerNorm(reduction_axes=[1, 2], feature_axes=-1).apply(variables, x) >>> np.testing.assert_allclose(y, y2, atol=1e-7) >>> y3 = nn.GroupNorm(num_groups=x.shape[-1]).apply(variables, x) >>> np.testing.assert_allclose(y, y3, atol=1e-7) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.epsilon"> epsilon<a class="headerlink" href="#flax.linen.InstanceNorm.epsilon" title="Permalink to this definition">#</a></dt> <dd>A small float added to variance to avoid dividing by zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.dtype"> dtype<a class="headerlink" href="#flax.linen.InstanceNorm.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the result (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.InstanceNorm.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.use_bias"> use_bias<a class="headerlink" href="#flax.linen.InstanceNorm.use_bias" title="Permalink to this definition">#</a></dt> <dd>If True, bias (beta) is added. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.use_scale"> use_scale<a class="headerlink" href="#flax.linen.InstanceNorm.use_scale" title="Permalink to this definition">#</a></dt> <dd>If True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.bias_init"> bias_init<a class="headerlink" href="#flax.linen.InstanceNorm.bias_init" title="Permalink to this definition">#</a></dt> <dd>Initializer for bias, by default, zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.scale_init"> scale_init<a class="headerlink" href="#flax.linen.InstanceNorm.scale_init" title="Permalink to this definition">#</a></dt> <dd>Initializer for scale, by default, one. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.feature_axes"> feature_axes<a class="headerlink" href="#flax.linen.InstanceNorm.feature_axes" title="Permalink to this definition">#</a></dt> <dd>Axes for features. The learned bias and scaling parameters will be in the shape defined by the feature axes. All other axes except the batch axes (which is assumed to be the leading axis) will be reduced. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[int, collections.abc.Sequence[int]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.axis_name"> axis_name<a class="headerlink" href="#flax.linen.InstanceNorm.axis_name" title="Permalink to this definition">#</a></dt> <dd>the axis name used to combine batch statistics from multiple devices. See <code class="docutils literal notranslate">jax.pmap</code> for a description of axis names (default: None). This is only needed if the model is subdivided across devices, i.e. the array being normalized is sharded across devices within a pmap or shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">str | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.axis_index_groups"> axis_index_groups<a class="headerlink" href="#flax.linen.InstanceNorm.axis_index_groups" title="Permalink to this definition">#</a></dt> <dd>groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, <code class="docutils literal notranslate">[[0, 1], [2, 3]]</code> would independently batch-normalize over the examples on the first two and last two devices. See <code class="docutils literal notranslate">jax.lax.psum</code> for more details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Any </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.use_fast_variance"> use_fast_variance<a class="headerlink" href="#flax.linen.InstanceNorm.use_fast_variance" title="Permalink to this definition">#</a></dt> <dd>If true, use a faster, but less numerically stable, calculation for the variance. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.InstanceNorm.__call__"> __call__(x, *, mask=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#InstanceNorm.__call__">[source]</a><a class="headerlink" href="#flax.linen.InstanceNorm.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies instance normalization on the input. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>x – the inputs</li> <li>mask – Binary array of shape broadcastable to <code class="docutils literal notranslate">inputs</code> tensor, indicating the positions for which the mean and variance should be computed.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">Normalized inputs (the same shape as inputs). </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.SpectralNorm"> class flax.linen.SpectralNorm(layer_instance, n_steps=1, epsilon=1e-12, dtype=None, param_dtype=<class 'jax.numpy.float32'>, error_on_non_matrix=False, collection_name='batch_stats', parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#SpectralNorm">[source]</a><a class="headerlink" href="#flax.linen.SpectralNorm" title="Permalink to this definition">#</a></dt> <dd>Spectral normalization. See: <ul class="simple"> <li><a class="reference external" href="https://arxiv.org/abs/1802.05957">https://arxiv.org/abs/1802.05957</a></li> <li><a class="reference external" href="https://arxiv.org/abs/1805.08318">https://arxiv.org/abs/1805.08318</a></li> <li><a class="reference external" href="https://arxiv.org/abs/1809.11096">https://arxiv.org/abs/1809.11096</a></li> </ul> Spectral normalization normalizes the weight params so that the spectral norm of the matrix is equal to 1. This is implemented as a layer wrapper where each wrapped layer will have its params spectral normalized before computing its <code class="docutils literal notranslate">__call__</code> output. <div class="admonition note"> Note The initialized variables dict will contain, in addition to a ‘params’ collection, a separate ‘batch_stats’ collection that will contain a <code class="docutils literal notranslate">u</code> vector and <code class="docutils literal notranslate">sigma</code> value, which are intermediate values used when performing spectral normalization. During training, we pass in <code class="docutils literal notranslate">update_stats=True</code> and <code class="docutils literal notranslate">mutable=['batch_stats']</code> so that <code class="docutils literal notranslate">u</code> and <code class="docutils literal notranslate">sigma</code> are updated with the most recently computed values using power iteration. This will help the power iteration method approximate the true singular value more accurately over time. During eval, we pass in <code class="docutils literal notranslate">update_stats=False</code> to ensure we get deterministic behavior from the model. </div> Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax, flax.linen as nn >>> import jax, jax.numpy as jnp >>> import optax >>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x, train): ... x = nn.Dense(3)(x) ... # only spectral normalize the params of the second Dense layer ... x = nn.SpectralNorm(nn.Dense(4))(x, update_stats=train) ... x = nn.Dense(5)(x) ... return x >>> # init >>> x = jnp.ones((1, 2)) >>> y = jnp.ones((1, 5)) >>> model = Foo() >>> variables = model.init(jax.random.PRNGKey(0), x, train=False) >>> flax.core.freeze(jax.tree_util.tree_map(jnp.shape, variables)) FrozenDict({ batch_stats: { SpectralNorm_0: { Dense_1/kernel/sigma: (), Dense_1/kernel/u: (1, 4), }, }, params: { Dense_0: { bias: (3,), kernel: (2, 3), }, Dense_1: { bias: (4,), kernel: (3, 4), }, Dense_2: { bias: (5,), kernel: (4, 5), }, }, }) >>> # train >>> def train_step(variables, x, y): ... def loss_fn(params): ... logits, updates = model.apply( ... {'params': params, 'batch_stats': variables['batch_stats']}, ... x, ... train=True, ... mutable=['batch_stats'], ... ) ... loss = jnp.mean(optax.l2_loss(predictions=logits, targets=y)) ... return loss, updates ... ... (loss, updates), grads = jax.value_and_grad(loss_fn, has_aux=True)( ... variables['params'] ... ) ... return { ... 'params': jax.tree_util.tree_map( ... lambda p, g: p - 0.1 * g, variables['params'], grads ... ), ... 'batch_stats': updates['batch_stats'], ... }, loss >>> for _ in range(10): ... variables, loss = train_step(variables, x, y) >>> # inference / eval >>> out = model.apply(variables, x, train=False) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SpectralNorm.layer_instance"> layer_instance<a class="headerlink" href="#flax.linen.SpectralNorm.layer_instance" title="Permalink to this definition">#</a></dt> <dd>Module instance that is wrapped with SpectralNorm <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd"><a class="reference internal" href="module.html#flax.linen.Module" title="flax.linen.module.Module">flax.linen.module.Module</a> </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SpectralNorm.n_steps"> n_steps<a class="headerlink" href="#flax.linen.SpectralNorm.n_steps" title="Permalink to this definition">#</a></dt> <dd>How many steps of power iteration to perform to approximate the singular value of the weight params. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SpectralNorm.epsilon"> epsilon<a class="headerlink" href="#flax.linen.SpectralNorm.epsilon" title="Permalink to this definition">#</a></dt> <dd>A small float added to l2-normalization to avoid dividing by zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SpectralNorm.dtype"> dtype<a class="headerlink" href="#flax.linen.SpectralNorm.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the result (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SpectralNorm.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.SpectralNorm.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SpectralNorm.error_on_non_matrix"> error_on_non_matrix<a class="headerlink" href="#flax.linen.SpectralNorm.error_on_non_matrix" title="Permalink to this definition">#</a></dt> <dd>Spectral normalization is only defined on matrices. By default, this module will return scalars unchanged and flatten higher-order tensors in their leading dimensions. Setting this flag to True will instead throw an error if a weight tensor with dimension greater than 2 is used by the layer. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SpectralNorm.collection_name"> collection_name<a class="headerlink" href="#flax.linen.SpectralNorm.collection_name" title="Permalink to this definition">#</a></dt> <dd>Name of the collection to store intermediate values used when performing spectral normalization. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">str </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.SpectralNorm.__call__"> __call__(*args, update_stats, **kwargs)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#SpectralNorm.__call__">[source]</a><a class="headerlink" href="#flax.linen.SpectralNorm.__call__" title="Permalink to this definition">#</a></dt> <dd>Compute the largest singular value of the weights in <code class="docutils literal notranslate">self.layer_instance</code> using power iteration and normalize the weights using this value before computing the <code class="docutils literal notranslate">__call__</code> output. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>*args – positional arguments to be passed into the call method of the underlying layer instance in <code class="docutils literal notranslate">self.layer_instance</code>.</li> <li>update_stats – if True, update the internal <code class="docutils literal notranslate">u</code> vector and <code class="docutils literal notranslate">sigma</code> value after computing their updated values using power iteration. This will help the power iteration method approximate the true singular value more accurately over time.</li> <li>**kwargs – keyword arguments to be passed into the call method of the underlying layer instance in <code class="docutils literal notranslate">self.layer_instance</code>.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">Output of the layer using spectral normalized weights. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.WeightNorm"> class flax.linen.WeightNorm(layer_instance, epsilon=1e-12, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_scale=True, scale_init=<function ones>, feature_axes=-1, variable_filter=<factory>, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#WeightNorm">[source]</a><a class="headerlink" href="#flax.linen.WeightNorm" title="Permalink to this definition">#</a></dt> <dd>L2 weight normalization (<a class="reference external" href="https://arxiv.org/abs/1602.07868">https://arxiv.org/abs/1602.07868</a>). Weight normalization normalizes the weight params so that the l2-norm of the matrix is equal to 1. This is implemented as a layer wrapper where each wrapped layer will have its params l2-normalized before computing its <code class="docutils literal notranslate">__call__</code> output. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax, flax.linen as nn >>> import jax, jax.numpy as jnp >>> class Baz(nn.Module): ... @nn.compact ... def __call__(self, x): ... return nn.Dense(2)(x) >>> class Bar(nn.Module): ... @nn.compact ... def __call__(self, x): ... x = Baz()(x) ... x = nn.Dense(3)(x) ... x = Baz()(x) ... x = nn.Dense(3)(x) ... return x >>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x): ... x = nn.Dense(3)(x) ... # l2-normalize all params of the second Dense layer ... x = nn.WeightNorm(nn.Dense(4), variable_filter=None)(x) ... x = nn.Dense(5)(x) ... # l2-normalize all kernels in the Bar submodule and all params in ... # the Baz submodule ... x = nn.WeightNorm(Bar(), variable_filter={'kernel', 'Baz'})(x) ... return x >>> # init >>> x = jnp.ones((1, 2)) >>> model = Foo() >>> variables = model.init(jax.random.key(0), x) >>> flax.core.freeze(jax.tree_util.tree_map(jnp.shape, variables)) FrozenDict({ params: { Bar_0: { Baz_0: { Dense_0: { bias: (2,), kernel: (5, 2), }, }, Baz_1: { Dense_0: { bias: (2,), kernel: (3, 2), }, }, Dense_0: { bias: (3,), kernel: (2, 3), }, Dense_1: { bias: (3,), kernel: (2, 3), }, }, Dense_0: { bias: (3,), kernel: (2, 3), }, Dense_1: { bias: (4,), kernel: (3, 4), }, Dense_2: { bias: (5,), kernel: (4, 5), }, WeightNorm_0: { Dense_1/bias/scale: (4,), Dense_1/kernel/scale: (4,), }, WeightNorm_1: { Bar_0/Baz_0/Dense_0/bias/scale: (2,), Bar_0/Baz_0/Dense_0/kernel/scale: (2,), Bar_0/Baz_1/Dense_0/bias/scale: (2,), Bar_0/Baz_1/Dense_0/kernel/scale: (2,), Bar_0/Dense_0/kernel/scale: (3,), Bar_0/Dense_1/kernel/scale: (3,), }, }, }) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.WeightNorm.layer_instance"> layer_instance<a class="headerlink" href="#flax.linen.WeightNorm.layer_instance" title="Permalink to this definition">#</a></dt> <dd>Module instance that is wrapped with WeightNorm <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd"><a class="reference internal" href="module.html#flax.linen.Module" title="flax.linen.module.Module">flax.linen.module.Module</a> </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.WeightNorm.epsilon"> epsilon<a class="headerlink" href="#flax.linen.WeightNorm.epsilon" title="Permalink to this definition">#</a></dt> <dd>A small float added to l2-normalization to avoid dividing by zero. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.WeightNorm.dtype"> dtype<a class="headerlink" href="#flax.linen.WeightNorm.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the result (default: infer from input and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.WeightNorm.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.WeightNorm.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.WeightNorm.use_scale"> use_scale<a class="headerlink" href="#flax.linen.WeightNorm.use_scale" title="Permalink to this definition">#</a></dt> <dd>If True, creates a learnable variable <code class="docutils literal notranslate">scale</code> that is multiplied to the <code class="docutils literal notranslate">layer_instance</code> variables after l2-normalization. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.WeightNorm.scale_init"> scale_init<a class="headerlink" href="#flax.linen.WeightNorm.scale_init" title="Permalink to this definition">#</a></dt> <dd>Initialization function for the scaling function. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.WeightNorm.feature_axes"> feature_axes<a class="headerlink" href="#flax.linen.WeightNorm.feature_axes" title="Permalink to this definition">#</a></dt> <dd>The feature axes dimension(s). The l2-norm is calculated by reducing the <code class="docutils literal notranslate">layer_instance</code> variables over the remaining (non-feature) axes. Therefore a separate l2-norm value is calculated and a separate scale (if <code class="docutils literal notranslate">use_scale=True</code>) is learned for each specified feature. By default, the trailing dimension is treated as the feature axis. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[int, collections.abc.Sequence[int]]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.WeightNorm.variable_filter"> variable_filter<a class="headerlink" href="#flax.linen.WeightNorm.variable_filter" title="Permalink to this definition">#</a></dt> <dd>An optional iterable that contains string items. The WeightNorm layer will selectively apply l2-normalization to the <code class="docutils literal notranslate">layer_instance</code> variables whose key path (delimited by ‘/’) has a match with <code class="docutils literal notranslate">variable_filter</code>. For example, <code class="docutils literal notranslate">variable_filter={'kernel'}</code> will only apply l2-normalization to variables whose key path contains ‘kernel’. By default, <code class="docutils literal notranslate">variable_filter={'kernel'}</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Iterable | None </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.WeightNorm.__call__"> __call__(*args, **kwargs)<a class="reference internal" href="../../_modules/flax/linen/normalization.html#WeightNorm.__call__">[source]</a><a class="headerlink" href="#flax.linen.WeightNorm.__call__" title="Permalink to this definition">#</a></dt> <dd>Compute the l2-norm of the weights in <code class="docutils literal notranslate">self.layer_instance</code> and normalize the weights using this value before computing the <code class="docutils literal notranslate">__call__</code> output. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>*args – positional arguments to be passed into the call method of the underlying layer instance in <code class="docutils literal notranslate">self.layer_instance</code>.</li> <li>**kwargs – keyword arguments to be passed into the call method of the underlying layer instance in <code class="docutils literal notranslate">self.layer_instance</code>.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">Output of the layer using l2-normalized weights. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> </div> <div class="section" id="combinators"> <h2>Combinators<a class="headerlink" href="#combinators" title="Permalink to this heading">#</a></h2> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.Sequential"> class flax.linen.Sequential(layers, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/combinators.html#Sequential">[source]</a><a class="headerlink" href="#flax.linen.Sequential" title="Permalink to this definition">#</a></dt> <dd>Applies a linear chain of Modules. Meant to be used only for the simple case of fusing together callables where the input of a particular module/op is the output of the previous one. Modules will be applied in the order that they are passed in the constructor. The <code class="docutils literal notranslate">__call__</code> method of Sequential accepts any input and forwards it to the first module it contains. It chains the output sequentially to the input of the next module and returns the output of the final module. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> class Foo(nn.Module): ... @nn.compact ... def __call__(self, x): ... return nn.Sequential([nn.Dense(4), ... nn.relu, ... nn.Dense(2), ... nn.log_softmax])(x) </pre></div> </div> Since <cite>Sequential.__call__</cite> is a <cite>compact</cite> method, you can also pass functions that construct Modules inline if you need shape inference: <div class="highlight-default notranslate"><div class="highlight"><pre>module = nn.Sequential([ # << more layers lambda x: SomeModule(x.shape[-1])(x), # shape inference # << more layers ]) </pre></div> </div> This combinator supports also layers that return multiple outputs if returned as a tuple or a dictionary. If the output of a layer is a <code class="docutils literal notranslate">tuple</code> it will be expanded as <code class="docutils literal notranslate">*args</code> in the next layer, if its a <code class="docutils literal notranslate">dict</code> it will be expanded as <code class="docutils literal notranslate">**kwargs</code>. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> class CrossAttentionBlock(nn.Module): ... num_heads: int = 2 ... qkv_features: int = 16 ... ... @nn.compact ... def __call__(self, query, key_value): ... output = nn.MultiHeadDotProductAttention( ... num_heads=self.num_heads, qkv_features=self.qkv_features)(query, ... key_value) ... output = nn.Dense(self.qkv_features)(output) ... return dict(query=output, key_value=key_value) # also works for tuples >>> from typing import Sequence >>> class CrossAttentionNetwork(nn.Module): ... num_layers: Sequence[int] ... ... @nn.compact ... def __call__(self, x): ... return nn.Sequential([CrossAttentionBlock() for _ in ... range(self.num_layers)])(query, key_value) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Sequential.layers"> layers<a class="headerlink" href="#flax.linen.Sequential.layers" title="Permalink to this definition">#</a></dt> <dd>A sequence of callables to be applied in order. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Sequence[collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="field-list simple"> <dt class="field-odd">Raises</dt> <dd class="field-odd">ValueError – If layers is not a sequence. </dd> </dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Sequential.__call__"> __call__(*args, **kwargs)<a class="reference internal" href="../../_modules/flax/linen/combinators.html#Sequential.__call__">[source]</a><a class="headerlink" href="#flax.linen.Sequential.__call__" title="Permalink to this definition">#</a></dt> <dd>Call self as a function. </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> </div> <div class="section" id="stochastic"> <h2>Stochastic<a class="headerlink" href="#stochastic" title="Permalink to this heading">#</a></h2> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.Dropout"> class flax.linen.Dropout(rate, broadcast_dims=(), deterministic=None, rng_collection='dropout', parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/stochastic.html#Dropout">[source]</a><a class="headerlink" href="#flax.linen.Dropout" title="Permalink to this definition">#</a></dt> <dd>Create a dropout layer. <div class="admonition note"> Note When using <a class="reference internal" href="module.html#flax.linen.Module.apply" title="flax.linen.Module.apply"><code class="xref py py-meth docutils literal notranslate">Module.apply()</code></a>, make sure to include an RNG seed named <code class="docutils literal notranslate">'dropout'</code>. Dropout isn’t necessary for variable initialization. </div> Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> class MLP(nn.Module): ... @nn.compact ... def __call__(self, x, train): ... x = nn.Dense(4)(x) ... x = nn.Dropout(0.5, deterministic=not train)(x) ... return x >>> model = MLP() >>> x = jnp.ones((1, 3)) >>> variables = model.init(jax.random.key(0), x, train=False) # don't use dropout >>> model.apply(variables, x, train=False) # don't use dropout Array([[-0.88686204, -0.5928178 , -0.5184689 , -0.4345976 ]], dtype=float32) >>> model.apply(variables, x, train=True, rngs={'dropout': jax.random.key(1)}) # use dropout Array([[ 0. , -1.1856356, -1.0369378, 0. ]], dtype=float32) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dropout.rate"> rate<a class="headerlink" href="#flax.linen.Dropout.rate" title="Permalink to this definition">#</a></dt> <dd>the dropout probability. (_not_ the keep rate!) <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dropout.broadcast_dims"> broadcast_dims<a class="headerlink" href="#flax.linen.Dropout.broadcast_dims" title="Permalink to this definition">#</a></dt> <dd>dimensions that will share the same dropout mask <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dropout.deterministic"> deterministic<a class="headerlink" href="#flax.linen.Dropout.deterministic" title="Permalink to this definition">#</a></dt> <dd>if false the inputs are scaled by <code class="docutils literal notranslate">1 / (1 - rate)</code> and masked, whereas if true, no mask is applied and the inputs are returned as is. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.Dropout.rng_collection"> rng_collection<a class="headerlink" href="#flax.linen.Dropout.rng_collection" title="Permalink to this definition">#</a></dt> <dd>the rng collection name to use when requesting an rng key. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">str </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Dropout.__call__"> __call__(inputs, deterministic=None, rng=None)<a class="reference internal" href="../../_modules/flax/linen/stochastic.html#Dropout.__call__">[source]</a><a class="headerlink" href="#flax.linen.Dropout.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies a random dropout mask to the input. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>inputs – the inputs that should be randomly masked.</li> <li>deterministic – if false the inputs are scaled by <code class="docutils literal notranslate">1 / (1 - rate)</code> and masked, whereas if true, no mask is applied and the inputs are returned as is.</li> <li>rng – an optional PRNGKey used as the random key, if not specified, one will be generated using <code class="docutils literal notranslate">make_rng</code> with the <code class="docutils literal notranslate">rng_collection</code> name.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">The masked inputs reweighted to preserve mean. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> </div> <div class="section" id="attention"> <h2>Attention<a class="headerlink" href="#attention" title="Permalink to this heading">#</a></h2> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention"> class flax.linen.MultiHeadDotProductAttention(num_heads, dtype=None, param_dtype=<class 'jax.numpy.float32'>, qkv_features=None, out_features=None, broadcast_dropout=True, dropout_rate=0.0, deterministic=None, precision=None, kernel_init=<function variance_scaling.<locals>.init>, out_kernel_init=None, bias_init=<function zeros>, out_bias_init=None, use_bias=True, attention_fn=<function dot_product_attention>, decode=False, normalize_qk=False, force_fp32_for_softmax=False, qkv_dot_general=None, out_dot_general=None, qkv_dot_general_cls=None, out_dot_general_cls=None, qk_attn_weights_einsum_cls=None, attn_weights_value_einsum_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/attention.html#MultiHeadDotProductAttention">[source]</a><a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention" title="Permalink to this definition">#</a></dt> <dd>Multi-head dot-product attention. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax >>> layer = nn.MultiHeadDotProductAttention(num_heads=8, qkv_features=16) >>> key1, key2, key3, key4, key5, key6 = jax.random.split(jax.random.key(0), 6) >>> shape = (4, 3, 2, 5) >>> q, k, v = jax.random.uniform(key1, shape), jax.random.uniform(key2, shape), jax.random.uniform(key3, shape) >>> variables = layer.init(jax.random.key(0), q) >>> # different inputs for inputs_q, inputs_k and inputs_v >>> out = layer.apply(variables, q, k, v) >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=k, inputs_v=k) >>> out = layer.apply(variables, q, k) >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=q) and layer.apply(variables, inputs_q=q, inputs_k=q, inputs_v=q) >>> out = layer.apply(variables, q) >>> attention_kwargs = dict( ... num_heads=8, ... qkv_features=16, ... kernel_init=nn.initializers.ones, ... bias_init=nn.initializers.zeros, ... dropout_rate=0.5, ... deterministic=False, ... ) >>> class Module(nn.Module): ... attention_kwargs: dict ... ... @nn.compact ... def __call__(self, x, dropout_rng=None): ... out1 = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng) ... out2 = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng) ... return out1, out2 >>> module = Module(attention_kwargs) >>> variables = module.init({'params': key1, 'dropout': key2}, q) >>> # out1 and out2 are different. >>> out1, out2 = module.apply(variables, q, rngs={'dropout': key3}) >>> # out3 and out4 are different. >>> # out1 and out3 are different. out2 and out4 are different. >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key4}) >>> # out1 and out2 are the same. >>> out1, out2 = module.apply(variables, q, dropout_rng=key5) >>> # out1 and out2 are the same as out3 and out4. >>> # providing a `dropout_rng` arg will take precedence over the `rngs` arg in `.apply` >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key6}, dropout_rng=key5) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.num_heads"> num_heads<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.num_heads" title="Permalink to this definition">#</a></dt> <dd>Number of attention heads. Features (i.e. inputs_q.shape[-1]) should be divisible by the number of heads. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.dtype"> dtype<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.dtype" title="Permalink to this definition">#</a></dt> <dd>The dtype of the computation (default: infer from inputs and params) <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.param_dtype" title="Permalink to this definition">#</a></dt> <dd>The dtype passed to parameter initializers (default: float32) <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.qkv_features"> qkv_features<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.qkv_features" title="Permalink to this definition">#</a></dt> <dd>Dimension of the key, query, and value. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.out_features"> out_features<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.out_features" title="Permalink to this definition">#</a></dt> <dd>Dimension of the last projection <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.broadcast_dropout"> broadcast_dropout<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.broadcast_dropout" title="Permalink to this definition">#</a></dt> <dd>Use a broadcasted dropout along batch dims. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.dropout_rate"> dropout_rate<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.dropout_rate" title="Permalink to this definition">#</a></dt> <dd>Dropout rate. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.deterministic"> deterministic<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.deterministic" title="Permalink to this definition">#</a></dt> <dd>If False, the attention weight is masked randomly using dropout, whereas if True, the attention weights are deterministic. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.precision"> precision<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.precision" title="Permalink to this definition">#</a></dt> <dd>Numerical precision of the computation see <code class="docutils literal notranslate">jax.lax.Precision</code> for details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[None, str, jax._src.lax.lax.Precision, tuple[str, str], tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.kernel_init" title="Permalink to this definition">#</a></dt> <dd>Initializer for the kernel of the Dense layers. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.out_kernel_init"> out_kernel_init<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.out_kernel_init" title="Permalink to this definition">#</a></dt> <dd>Optional Initializer for the kernel of the output Dense layer, if None, <code class="docutils literal notranslate">kernel_init</code> will be used. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.bias_init"> bias_init<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.bias_init" title="Permalink to this definition">#</a></dt> <dd>Initializer for the bias of the Dense layers. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.out_bias_init"> out_bias_init<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.out_bias_init" title="Permalink to this definition">#</a></dt> <dd>Optional Initializer for the bias of the output Dense layer, if None, <code class="docutils literal notranslate">bias_init</code> will be used. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.use_bias"> use_bias<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.use_bias" title="Permalink to this definition">#</a></dt> <dd>Whether pointwise QKVO dense transforms use bias. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.attention_fn"> attention_fn<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.attention_fn" title="Permalink to this definition">#</a></dt> <dd>dot_product_attention or compatible function. Accepts query, key, value, and returns output of shape <code class="docutils literal notranslate">[bs, dim1, dim2, ..., dimN,, num_heads, value_channels]</code> <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Union[jax.Array, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.decode"> decode<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.decode" title="Permalink to this definition">#</a></dt> <dd>Whether to prepare and use an autoregressive cache. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.normalize_qk"> normalize_qk<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.normalize_qk" title="Permalink to this definition">#</a></dt> <dd>Should QK normalization be applied (arxiv.org/abs/2302.05442). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.qk_attn_weights_einsum_cls"> qk_attn_weights_einsum_cls<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.qk_attn_weights_einsum_cls" title="Permalink to this definition">#</a></dt> <dd>factory function to create the einsum for computing the attention weights. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], collections.abc.Callable[[…], Union[jax.Array, Any]]] | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.attn_weights_value_einsum_cls"> attn_weights_value_einsum_cls<a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.attn_weights_value_einsum_cls" title="Permalink to this definition">#</a></dt> <dd>factory function to create the einsum for computing the product of the attention weights and the values. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], collections.abc.Callable[[…], Union[jax.Array, Any]]] | None </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.MultiHeadDotProductAttention.__call__"> __call__(inputs_q, inputs_k=None, inputs_v=None, *, inputs_kv=None, mask=None, deterministic=None, dropout_rng=None, sow_weights=False)<a class="reference internal" href="../../_modules/flax/linen/attention.html#MultiHeadDotProductAttention.__call__">[source]</a><a class="headerlink" href="#flax.linen.MultiHeadDotProductAttention.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies multi-head dot product attention on the input data. Projects the inputs into multi-headed query, key, and value vectors, applies dot-product attention and project the results to an output vector. If both inputs_k and inputs_v are None, they will both copy the value of inputs_q (self attention). If only inputs_v is None, it will copy the value of inputs_k. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>inputs_q – input queries of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>.</li> <li>inputs_k – key of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>. If None, inputs_k will copy the value of inputs_q.</li> <li>inputs_v – values of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>. If None, inputs_v will copy the value of inputs_k.</li> <li>inputs_kv – key/values of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>. If None, inputs_kv will copy the value of inputs_q. This arg will be deprecated soon. Use inputs_k and inputs_v instead.</li> <li>mask – attention mask of shape <code class="docutils literal notranslate">[batch_sizes..., num_heads, query_length, key/value_length]</code>. Attention weights are masked out if their corresponding mask value is <code class="docutils literal notranslate">False</code>.</li> <li>deterministic – if false, the attention weight is masked randomly using dropout, whereas if true, the attention weights are deterministic.</li> <li>dropout_rng – optional rng key to pass to the attention layer’s dropout mask. Otherwise, self.make_rng(‘dropout’) is used instead.</li> <li>sow_weights – if <code class="docutils literal notranslate">True</code>, the attention weights are sowed into the ‘intermediates’ collection. Remember to mark ‘intermediates’ as mutable via <code class="docutils literal notranslate">mutable=['intermediates']</code> in order to have that collection returned.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">output of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention"> class flax.linen.MultiHeadAttention(num_heads, dtype=None, param_dtype=<class 'jax.numpy.float32'>, qkv_features=None, out_features=None, broadcast_dropout=True, dropout_rate=0.0, deterministic=None, precision=None, kernel_init=<function variance_scaling.<locals>.init>, out_kernel_init=None, bias_init=<function zeros>, out_bias_init=None, use_bias=True, attention_fn=<function dot_product_attention>, decode=False, normalize_qk=False, force_fp32_for_softmax=False, qkv_dot_general=None, out_dot_general=None, qkv_dot_general_cls=None, out_dot_general_cls=None, qk_attn_weights_einsum_cls=None, attn_weights_value_einsum_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/attention.html#MultiHeadAttention">[source]</a><a class="headerlink" href="#flax.linen.MultiHeadAttention" title="Permalink to this definition">#</a></dt> <dd>Multi-head dot-product attention. Alias for <code class="docutils literal notranslate">MultiHeadDotProductAttention</code>. NOTE: <code class="docutils literal notranslate">MultiHeadAttention</code> is a wrapper of <code class="docutils literal notranslate">MultiHeadDotProductAttention</code>, and so their implementations are identical. However <code class="docutils literal notranslate">MultiHeadAttention</code> layers will, by default, be named <code class="docutils literal notranslate">MultiHeadAttention_{index}</code>, whereas <code class="docutils literal notranslate">MultiHeadDotProductAttention</code> will be named <code class="docutils literal notranslate">MultiHeadDotProductAttention_{index}</code>. Therefore, this could affect checkpointing, param collection names and RNG threading (since the layer name is used when generating new RNG’s) within the module. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax >>> layer = nn.MultiHeadAttention(num_heads=8, qkv_features=16) >>> key1, key2, key3, key4, key5, key6 = jax.random.split(jax.random.key(0), 6) >>> shape = (4, 3, 2, 5) >>> q, k, v = jax.random.uniform(key1, shape), jax.random.uniform(key2, shape), jax.random.uniform(key3, shape) >>> variables = layer.init(jax.random.key(0), q) >>> # different inputs for inputs_q, inputs_k and inputs_v >>> out = layer.apply(variables, q, k, v) >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=k, inputs_v=k) >>> out = layer.apply(variables, q, k) >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=q) and layer.apply(variables, inputs_q=q, inputs_k=q, inputs_v=q) >>> out = layer.apply(variables, q) >>> attention_kwargs = dict( ... num_heads=8, ... qkv_features=16, ... kernel_init=nn.initializers.ones, ... bias_init=nn.initializers.zeros, ... dropout_rate=0.5, ... deterministic=False, ... ) >>> class Module(nn.Module): ... attention_kwargs: dict ... ... @nn.compact ... def __call__(self, x, dropout_rng=None): ... out1 = nn.MultiHeadAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng) ... out2 = nn.MultiHeadAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng) ... return out1, out2 >>> module = Module(attention_kwargs) >>> variables = module.init({'params': key1, 'dropout': key2}, q) >>> # out1 and out2 are different. >>> out1, out2 = module.apply(variables, q, rngs={'dropout': key3}) >>> # out3 and out4 are different. >>> # out1 and out3 are different. out2 and out4 are different. >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key4}) >>> # out1 and out2 are the same. >>> out1, out2 = module.apply(variables, q, dropout_rng=key5) >>> # out1 and out2 are the same as out3 and out4. >>> # providing a `dropout_rng` arg will take precedence over the `rngs` arg in `.apply` >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key6}, dropout_rng=key5) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.num_heads"> num_heads<a class="headerlink" href="#flax.linen.MultiHeadAttention.num_heads" title="Permalink to this definition">#</a></dt> <dd>number of attention heads. Features (i.e. inputs_q.shape[-1]) should be divisible by the number of heads. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.dtype"> dtype<a class="headerlink" href="#flax.linen.MultiHeadAttention.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: infer from inputs and params) <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.MultiHeadAttention.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32) <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.qkv_features"> qkv_features<a class="headerlink" href="#flax.linen.MultiHeadAttention.qkv_features" title="Permalink to this definition">#</a></dt> <dd>dimension of the key, query, and value. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.out_features"> out_features<a class="headerlink" href="#flax.linen.MultiHeadAttention.out_features" title="Permalink to this definition">#</a></dt> <dd>dimension of the last projection <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.broadcast_dropout"> broadcast_dropout<a class="headerlink" href="#flax.linen.MultiHeadAttention.broadcast_dropout" title="Permalink to this definition">#</a></dt> <dd>bool: use a broadcasted dropout along batch dims. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.dropout_rate"> dropout_rate<a class="headerlink" href="#flax.linen.MultiHeadAttention.dropout_rate" title="Permalink to this definition">#</a></dt> <dd>dropout rate <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">float </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.deterministic"> deterministic<a class="headerlink" href="#flax.linen.MultiHeadAttention.deterministic" title="Permalink to this definition">#</a></dt> <dd>if false, the attention weight is masked randomly using dropout, whereas if true, the attention weights are deterministic. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.precision"> precision<a class="headerlink" href="#flax.linen.MultiHeadAttention.precision" title="Permalink to this definition">#</a></dt> <dd>numerical precision of the computation see <code class="docutils literal notranslate">jax.lax.Precision</code> for details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[None, str, jax._src.lax.lax.Precision, tuple[str, str], tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.MultiHeadAttention.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the kernel of the Dense layers. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.bias_init"> bias_init<a class="headerlink" href="#flax.linen.MultiHeadAttention.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias of the Dense layers. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.use_bias"> use_bias<a class="headerlink" href="#flax.linen.MultiHeadAttention.use_bias" title="Permalink to this definition">#</a></dt> <dd>bool: whether pointwise QKVO dense transforms use bias. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.attention_fn"> attention_fn<a class="headerlink" href="#flax.linen.MultiHeadAttention.attention_fn" title="Permalink to this definition">#</a></dt> <dd>dot_product_attention or compatible function. Accepts query, key, value, and returns output of shape <code class="docutils literal notranslate">[bs, dim1, dim2, ..., dimN,, num_heads, value_channels]</code> <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Union[jax.Array, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.decode"> decode<a class="headerlink" href="#flax.linen.MultiHeadAttention.decode" title="Permalink to this definition">#</a></dt> <dd>whether to prepare and use an autoregressive cache. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.normalize_qk"> normalize_qk<a class="headerlink" href="#flax.linen.MultiHeadAttention.normalize_qk" title="Permalink to this definition">#</a></dt> <dd>should QK normalization be applied (arxiv.org/abs/2302.05442). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.MultiHeadAttention.__call__"> __call__(inputs_q, inputs_k=None, inputs_v=None, *, inputs_kv=None, mask=None, deterministic=None, dropout_rng=None, sow_weights=False)<a class="headerlink" href="#flax.linen.MultiHeadAttention.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies multi-head dot product attention on the input data. Projects the inputs into multi-headed query, key, and value vectors, applies dot-product attention and project the results to an output vector. If both inputs_k and inputs_v are None, they will both copy the value of inputs_q (self attention). If only inputs_v is None, it will copy the value of inputs_k. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>inputs_q – input queries of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>.</li> <li>inputs_k – key of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>. If None, inputs_k will copy the value of inputs_q.</li> <li>inputs_v – values of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>. If None, inputs_v will copy the value of inputs_k.</li> <li>inputs_kv – key/values of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>. If None, inputs_kv will copy the value of inputs_q. This arg will be deprecated soon. Use inputs_k and inputs_v instead.</li> <li>mask – attention mask of shape <code class="docutils literal notranslate">[batch_sizes..., num_heads, query_length, key/value_length]</code>. Attention weights are masked out if their corresponding mask value is <code class="docutils literal notranslate">False</code>.</li> <li>deterministic – if false, the attention weight is masked randomly using dropout, whereas if true, the attention weights are deterministic.</li> <li>dropout_rng – optional rng key to pass to the attention layer’s dropout mask. Otherwise, self.make_rng(‘dropout’) is used instead.</li> <li>sow_weights – if <code class="docutils literal notranslate">True</code>, the attention weights are sowed into the ‘intermediates’ collection. Remember to mark ‘intermediates’ as mutable via <code class="docutils literal notranslate">mutable=['intermediates']</code> in order to have that collection returned.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">output of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.SelfAttention"> class flax.linen.SelfAttention(num_heads, dtype=None, param_dtype=<class 'jax.numpy.float32'>, qkv_features=None, out_features=None, broadcast_dropout=True, dropout_rate=0.0, deterministic=None, precision=None, kernel_init=<function variance_scaling.<locals>.init>, out_kernel_init=None, bias_init=<function zeros>, out_bias_init=None, use_bias=True, attention_fn=<function dot_product_attention>, decode=False, normalize_qk=False, force_fp32_for_softmax=False, qkv_dot_general=None, out_dot_general=None, qkv_dot_general_cls=None, out_dot_general_cls=None, qk_attn_weights_einsum_cls=None, attn_weights_value_einsum_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/attention.html#SelfAttention">[source]</a><a class="headerlink" href="#flax.linen.SelfAttention" title="Permalink to this definition">#</a></dt> <dd>Self-attention special case of multi-head dot-product attention. This layer is deprecated in favor of <code class="docutils literal notranslate">MultiHeadDotProductAttention</code>. <dl> <dt>Example usage::</dt><dd><div class="doctest highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> layer = nn.MultiHeadDotProductAttention(num_heads=8, qkv_features=16) >>> variables = layer.init(jax.random.key(0), jnp.ones((4, 3, 2, 5))) </pre></div> </div> </dd> </dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.SelfAttention.__call__"> __call__(inputs_q, mask=None, deterministic=None, dropout_rng=None, sow_weights=False)<a class="reference internal" href="../../_modules/flax/linen/attention.html#SelfAttention.__call__">[source]</a><a class="headerlink" href="#flax.linen.SelfAttention.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies multi-head dot product self-attention on the input data. Projects the inputs into multi-headed query, key, and value vectors, applies dot-product attention and project the results to an output vector. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>inputs_q – input queries of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>.</li> <li>mask – attention mask of shape <code class="docutils literal notranslate">[batch_sizes..., num_heads, query_length, key/value_length]</code>. Attention weights are masked out if their corresponding mask value is <code class="docutils literal notranslate">False</code>.</li> <li>deterministic – if false, the attention weight is masked randomly using dropout, whereas if true, the attention weights are deterministic.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">output of shape <code class="docutils literal notranslate">[batch_sizes..., length, features]</code>. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.dot_product_attention_weights"> flax.linen.dot_product_attention_weights(query, key, bias=None, mask=None, broadcast_dropout=True, dropout_rng=None, dropout_rate=0.0, deterministic=False, dtype=None, precision=None, module=None, force_fp32_for_softmax=False, einsum_dot_general=None, einsum=None)<a class="reference internal" href="../../_modules/flax/linen/attention.html#dot_product_attention_weights">[source]</a><a class="headerlink" href="#flax.linen.dot_product_attention_weights" title="Permalink to this definition">#</a></dt> <dd>Computes dot-product attention weights given query and key. Used by <a class="reference internal" href="#flax.linen.dot_product_attention" title="flax.linen.dot_product_attention"><code class="xref py py-func docutils literal notranslate">dot_product_attention()</code></a>, which is what you’ll most likely use. But if you want access to the attention weights for introspection, then you can directly call this function and call einsum yourself. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>query – queries for calculating attention with shape of <code class="docutils literal notranslate">[batch..., q_length, num_heads, qk_depth_per_head]</code>.</li> <li>key – keys for calculating attention with shape of <code class="docutils literal notranslate">[batch..., kv_length, num_heads, qk_depth_per_head]</code>.</li> <li>bias – bias for the attention weights. This should be broadcastable to the shape <code class="docutils literal notranslate">[batch..., num_heads, q_length, kv_length]</code>. This can be used for incorporating causal masks, padding masks, proximity bias, etc.</li> <li>mask – mask for the attention weights. This should be broadcastable to the shape <code class="docutils literal notranslate">[batch..., num_heads, q_length, kv_length]</code>. This can be used for incorporating causal masks. Attention weights are masked out if their corresponding mask value is <code class="docutils literal notranslate">False</code>.</li> <li>broadcast_dropout – bool: use a broadcasted dropout along batch dims.</li> <li>dropout_rng – JAX PRNGKey: to be used for dropout</li> <li>dropout_rate – dropout rate</li> <li>deterministic – bool, deterministic or not (to apply dropout)</li> <li>dtype – the dtype of the computation (default: infer from inputs and params)</li> <li>precision – numerical precision of the computation see <code class="docutils literal notranslate">jax.lax.Precision</code> for details.</li> <li>module – the Module that will sow the attention weights into the ‘intermediates’ collection. Remember to mark ‘intermediates’ as mutable via <code class="docutils literal notranslate">mutable=['intermediates']</code> in order to have that collection returned. If <code class="docutils literal notranslate">module</code> is None, the attention weights will not be sowed.</li> <li>force_fp32_for_softmax – bool, whether to force the softmax to be computed in fp32. This is useful for mixed-precision training where higher precision is desired for numerical stability.</li> <li>einsum_dot_general – the dot_general to use in einsum.</li> <li>einsum – If unspecified, default <cite>jnp.einsum</cite> will be used. This argument is mutually exclusive with <cite>precision</cite> and <cite>einsum_dot_general</cite>.</li> </ul> </dd> <dt class="field-even">Raises</dt> <dd class="field-even">ValueError – if both <cite>precision</cite>/<cite>einsum_dot_general</cite> and <cite>einsum</cite> are specified. </dd> <dt class="field-odd">Returns</dt> <dd class="field-odd">Output of shape <code class="docutils literal notranslate">[batch..., num_heads, q_length, kv_length]</code>. </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.dot_product_attention"> flax.linen.dot_product_attention(query, key, value, bias=None, mask=None, broadcast_dropout=True, dropout_rng=None, dropout_rate=0.0, deterministic=False, dtype=None, precision=None, module=None, force_fp32_for_softmax=False, einsum_dot_general=None, qk_attn_weights_einsum=None, attn_weights_value_einsum=None)<a class="reference internal" href="../../_modules/flax/linen/attention.html#dot_product_attention">[source]</a><a class="headerlink" href="#flax.linen.dot_product_attention" title="Permalink to this definition">#</a></dt> <dd>Computes dot-product attention given query, key, and value. This is the core function for applying attention based on <a class="reference external" href="https://arxiv.org/abs/1706.03762">https://arxiv.org/abs/1706.03762</a>. It calculates the attention weights given query and key and combines the values using the attention weights. <div class="admonition note"> Note <code class="docutils literal notranslate">query</code>, <code class="docutils literal notranslate">key</code>, <code class="docutils literal notranslate">value</code> needn’t have any batch dimensions. </div> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>query – queries for calculating attention with shape of <code class="docutils literal notranslate">[batch..., q_length, num_heads, qk_depth_per_head]</code>.</li> <li>key – keys for calculating attention with shape of <code class="docutils literal notranslate">[batch..., kv_length, num_heads, qk_depth_per_head]</code>.</li> <li>value – values to be used in attention with shape of <code class="docutils literal notranslate">[batch..., kv_length, num_heads, v_depth_per_head]</code>.</li> <li>bias – bias for the attention weights. This should be broadcastable to the shape <code class="docutils literal notranslate">[batch..., num_heads, q_length, kv_length]</code>. This can be used for incorporating causal masks, padding masks, proximity bias, etc.</li> <li>mask – mask for the attention weights. This should be broadcastable to the shape <code class="docutils literal notranslate">[batch..., num_heads, q_length, kv_length]</code>. This can be used for incorporating causal masks. Attention weights are masked out if their corresponding mask value is <code class="docutils literal notranslate">False</code>.</li> <li>broadcast_dropout – bool: use a broadcasted dropout along batch dims.</li> <li>dropout_rng – JAX PRNGKey: to be used for dropout</li> <li>dropout_rate – dropout rate</li> <li>deterministic – bool, deterministic or not (to apply dropout)</li> <li>dtype – the dtype of the computation (default: infer from inputs)</li> <li>precision – numerical precision of the computation see <a href="#id3">``</a>jax.lax.Precision` for details.</li> <li>module – the Module that will sow the attention weights into the ‘intermediates’ collection. Remember to mark ‘intermediates’ as mutable via <code class="docutils literal notranslate">mutable=['intermediates']</code> in order to have that collection returned. If <code class="docutils literal notranslate">module</code> is None, the attention weights will not be sowed.</li> <li>force_fp32_for_softmax – bool, whether to force the softmax to be computed in fp32. This is useful for mixed-precision training where higher precision is desired for numerical stability.</li> <li>einsum_dot_general – the dot_general to use in <cite>jnp.einsum</cite>.</li> <li>qk_attn_weights_einsum – the einsum for computing the attention weights. When unspecified, the default <cite>jnp.einsum</cite> will be used. This argument is mutually exclusive with <cite>precision</cite> and <cite>einsum_dot_general</cite>.</li> <li>attn_weights_value_einsum – the einsum for computing the product of the attention weights and the values. When unspecified, the default <cite>jnp.einsum</cite> will be used. This argument is mutually exclusive with <cite>precision</cite> and <cite>einsum_dot_general</cite>.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">Output of shape <code class="docutils literal notranslate">[batch..., q_length, num_heads, v_depth_per_head]</code>. </dd> <dt class="field-odd">Raises</dt> <dd class="field-odd"><ul class="simple"> <li>ValueError – if both <cite>precision</cite>/<cite>einsum_dot_general</cite> and</li> <li>qk_attn_weights_einsum – specified.</li> </ul> </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.make_attention_mask"> flax.linen.make_attention_mask(query_input, key_input, pairwise_fn=<jnp.ufunc 'multiply'>, extra_batch_dims=0, dtype=<class 'jax.numpy.float32'>)<a class="reference internal" href="../../_modules/flax/linen/attention.html#make_attention_mask">[source]</a><a class="headerlink" href="#flax.linen.make_attention_mask" title="Permalink to this definition">#</a></dt> <dd>Mask-making helper for attention weights. In case of 1d inputs (i.e., <code class="docutils literal notranslate">[batch..., len_q]</code>, <code class="docutils literal notranslate">[batch..., len_kv]</code>, the attention weights will be <code class="docutils literal notranslate">[batch..., heads, len_q, len_kv]</code> and this function will produce <code class="docutils literal notranslate">[batch..., 1, len_q, len_kv]</code>. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>query_input – a batched, flat input of query_length size</li> <li>key_input – a batched, flat input of key_length size</li> <li>pairwise_fn – broadcasting elementwise comparison function</li> <li>extra_batch_dims – number of extra batch dims to add singleton axes for, none by default</li> <li>dtype – mask return dtype</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A <code class="docutils literal notranslate">[batch..., 1, len_q, len_kv]</code> shaped mask for 1d attention. </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.linen.make_causal_mask"> flax.linen.make_causal_mask(x, extra_batch_dims=0, dtype=<class 'jax.numpy.float32'>)<a class="reference internal" href="../../_modules/flax/linen/attention.html#make_causal_mask">[source]</a><a class="headerlink" href="#flax.linen.make_causal_mask" title="Permalink to this definition">#</a></dt> <dd>Make a causal mask for self-attention. In case of 1d inputs (i.e., <code class="docutils literal notranslate">[batch..., len]</code>, the self-attention weights will be <code class="docutils literal notranslate">[batch..., heads, len, len]</code> and this function will produce a causal mask of shape <code class="docutils literal notranslate">[batch..., 1, len, len]</code>. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>x – input array of shape <code class="docutils literal notranslate">[batch..., len]</code></li> <li>extra_batch_dims – number of batch dims to add singleton axes for, none by default</li> <li>dtype – mask return dtype</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A <code class="docutils literal notranslate">[batch..., 1, len, len]</code> shaped causal mask for 1d attention. </dd> </dl> </dd></dl> </div> <div class="section" id="recurrent"> <h2>Recurrent<a class="headerlink" href="#recurrent" title="Permalink to this heading">#</a></h2> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.RNNCellBase"> class flax.linen.RNNCellBase(parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#RNNCellBase">[source]</a><a class="headerlink" href="#flax.linen.RNNCellBase" title="Permalink to this definition">#</a></dt> <dd>RNN cell base class. <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.RNNCellBase.__call__"> __call__(**kwargs)<a class="headerlink" href="#flax.linen.RNNCellBase.__call__" title="Permalink to this definition">#</a></dt> <dd>Call self as a function. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.RNNCellBase.initialize_carry"> initialize_carry(rng, input_shape)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#RNNCellBase.initialize_carry">[source]</a><a class="headerlink" href="#flax.linen.RNNCellBase.initialize_carry" title="Permalink to this definition">#</a></dt> <dd>Initialize the RNN cell carry. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rng – random number generator passed to the init_fn.</li> <li>input_shape – a tuple providing the shape of the input to the cell.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">An initialized carry for the given RNN cell. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> <tr class="row-odd"><td><a class="reference internal" href="#flax.linen.RNNCellBase.initialize_carry" title="flax.linen.RNNCellBase.initialize_carry"><code class="xref py py-obj docutils literal notranslate">initialize_carry</code></a>(rng, input_shape)</td> <td>Initialize the RNN cell carry.</td> </tr> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.LSTMCell"> class flax.linen.LSTMCell(features, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#LSTMCell">[source]</a><a class="headerlink" href="#flax.linen.LSTMCell" title="Permalink to this definition">#</a></dt> <dd>LSTM cell. The mathematical definition of the cell is as follows <div class="math notranslate nohighlight"> \[\begin{split}\begin{array}{ll} i = \sigma(W_{ii} x + W_{hi} h + b_{hi}) \\ f = \sigma(W_{if} x + W_{hf} h + b_{hf}) \\ g = \tanh(W_{ig} x + W_{hg} h + b_{hg}) \\ o = \sigma(W_{io} x + W_{ho} h + b_{ho}) \\ c' = f * c + i * g \\ h' = o * \tanh(c') \\ \end{array}\end{split}\]</div> where x is the input, h is the output of the previous time step, and c is the memory. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> x = jax.random.normal(jax.random.key(0), (2, 3)) >>> layer = nn.LSTMCell(features=4) >>> carry = layer.initialize_carry(jax.random.key(1), x.shape) >>> variables = layer.init(jax.random.key(2), carry, x) >>> new_carry, out = layer.apply(variables, carry, x) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.features"> features<a class="headerlink" href="#flax.linen.LSTMCell.features" title="Permalink to this definition">#</a></dt> <dd>number of output features. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.gate_fn"> gate_fn<a class="headerlink" href="#flax.linen.LSTMCell.gate_fn" title="Permalink to this definition">#</a></dt> <dd>activation function used for gates (default: sigmoid). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.activation_fn"> activation_fn<a class="headerlink" href="#flax.linen.LSTMCell.activation_fn" title="Permalink to this definition">#</a></dt> <dd>activation function used for output and memory update (default: tanh). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.LSTMCell.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the input (default: lecun_normal). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.recurrent_kernel_init"> recurrent_kernel_init<a class="headerlink" href="#flax.linen.LSTMCell.recurrent_kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.bias_init"> bias_init<a class="headerlink" href="#flax.linen.LSTMCell.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias parameters (default: initializers.zeros_init()) <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.dtype"> dtype<a class="headerlink" href="#flax.linen.LSTMCell.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: infer from inputs and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.LSTMCell.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.__call__"> __call__(carry, inputs)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#LSTMCell.__call__">[source]</a><a class="headerlink" href="#flax.linen.LSTMCell.__call__" title="Permalink to this definition">#</a></dt> <dd>A long short-term memory (LSTM) cell. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>carry – the hidden state of the LSTM cell, initialized using <code class="docutils literal notranslate">LSTMCell.initialize_carry</code>.</li> <li>inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A tuple with the new carry and the output. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.LSTMCell.initialize_carry"> initialize_carry(rng, input_shape)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#LSTMCell.initialize_carry">[source]</a><a class="headerlink" href="#flax.linen.LSTMCell.initialize_carry" title="Permalink to this definition">#</a></dt> <dd>Initialize the RNN cell carry. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rng – random number generator passed to the init_fn.</li> <li>input_shape – a tuple providing the shape of the input to the cell.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">An initialized carry for the given RNN cell. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> <tr class="row-odd"><td><a class="reference internal" href="#flax.linen.LSTMCell.initialize_carry" title="flax.linen.LSTMCell.initialize_carry"><code class="xref py py-obj docutils literal notranslate">initialize_carry</code></a>(rng, input_shape)</td> <td>Initialize the RNN cell carry.</td> </tr> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell"> class flax.linen.OptimizedLSTMCell(features, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#OptimizedLSTMCell">[source]</a><a class="headerlink" href="#flax.linen.OptimizedLSTMCell" title="Permalink to this definition">#</a></dt> <dd>More efficient LSTM Cell that concatenates state components before matmul. The parameters are compatible with <code class="docutils literal notranslate">LSTMCell</code>. Note that this cell is often faster than <code class="docutils literal notranslate">LSTMCell</code> as long as the hidden size is roughly <= 2048 units. The mathematical definition of the cell is the same as <code class="docutils literal notranslate">LSTMCell</code> and as follows <div class="math notranslate nohighlight"> \[\begin{split}\begin{array}{ll} i = \sigma(W_{ii} x + W_{hi} h + b_{hi}) \\ f = \sigma(W_{if} x + W_{hf} h + b_{hf}) \\ g = \tanh(W_{ig} x + W_{hg} h + b_{hg}) \\ o = \sigma(W_{io} x + W_{ho} h + b_{ho}) \\ c' = f * c + i * g \\ h' = o * \tanh(c') \\ \end{array}\end{split}\]</div> where x is the input, h is the output of the previous time step, and c is the memory. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> x = jax.random.normal(jax.random.key(0), (2, 3)) >>> layer = nn.OptimizedLSTMCell(features=4) >>> carry = layer.initialize_carry(jax.random.key(1), x.shape) >>> variables = layer.init(jax.random.key(2), carry, x) >>> new_carry, out = layer.apply(variables, carry, x) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell.gate_fn"> gate_fn<a class="headerlink" href="#flax.linen.OptimizedLSTMCell.gate_fn" title="Permalink to this definition">#</a></dt> <dd>activation function used for gates (default: sigmoid). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell.activation_fn"> activation_fn<a class="headerlink" href="#flax.linen.OptimizedLSTMCell.activation_fn" title="Permalink to this definition">#</a></dt> <dd>activation function used for output and memory update (default: tanh). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.OptimizedLSTMCell.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the input (default: lecun_normal). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell.recurrent_kernel_init"> recurrent_kernel_init<a class="headerlink" href="#flax.linen.OptimizedLSTMCell.recurrent_kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell.bias_init"> bias_init<a class="headerlink" href="#flax.linen.OptimizedLSTMCell.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias parameters (default: initializers.zeros_init()). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell.dtype"> dtype<a class="headerlink" href="#flax.linen.OptimizedLSTMCell.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: infer from inputs and params). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.OptimizedLSTMCell.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell.__call__"> __call__(carry, inputs)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#OptimizedLSTMCell.__call__">[source]</a><a class="headerlink" href="#flax.linen.OptimizedLSTMCell.__call__" title="Permalink to this definition">#</a></dt> <dd>An optimized long short-term memory (LSTM) cell. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>carry – the hidden state of the LSTM cell, initialized using <code class="docutils literal notranslate">LSTMCell.initialize_carry</code>.</li> <li>inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A tuple with the new carry and the output. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.OptimizedLSTMCell.initialize_carry"> initialize_carry(rng, input_shape)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#OptimizedLSTMCell.initialize_carry">[source]</a><a class="headerlink" href="#flax.linen.OptimizedLSTMCell.initialize_carry" title="Permalink to this definition">#</a></dt> <dd>Initialize the RNN cell carry. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rng – random number generator passed to the init_fn.</li> <li>input_shape – a tuple providing the shape of the input to the cell.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">An initialized carry for the given RNN cell. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> <tr class="row-odd"><td><a class="reference internal" href="#flax.linen.OptimizedLSTMCell.initialize_carry" title="flax.linen.OptimizedLSTMCell.initialize_carry"><code class="xref py py-obj docutils literal notranslate">initialize_carry</code></a>(rng, input_shape)</td> <td>Initialize the RNN cell carry.</td> </tr> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell"> class flax.linen.ConvLSTMCell(features, kernel_size, strides=None, padding='SAME', use_bias=True, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#ConvLSTMCell">[source]</a><a class="headerlink" href="#flax.linen.ConvLSTMCell" title="Permalink to this definition">#</a></dt> <dd>A convolutional LSTM cell. The implementation is based on xingjian2015convolutional. Given x_t and the previous state (h_{t-1}, c_{t-1}) the core computes <div class="math notranslate nohighlight"> \[\begin{split}\begin{array}{ll} i_t = \sigma(W_{ii} * x_t + W_{hi} * h_{t-1} + b_i) \\ f_t = \sigma(W_{if} * x_t + W_{hf} * h_{t-1} + b_f) \\ g_t = \tanh(W_{ig} * x_t + W_{hg} * h_{t-1} + b_g) \\ o_t = \sigma(W_{io} * x_t + W_{ho} * h_{t-1} + b_o) \\ c_t = f_t c_{t-1} + i_t g_t \\ h_t = o_t \tanh(c_t) \end{array}\end{split}\]</div> where * denotes the convolution operator; i_t, f_t, o_t are input, forget and output gate activations, and g_t is a vector of cell updates. <div class="admonition note"> Note <dl class="simple"> <dt>Forget gate initialization:</dt><dd>Following jozefowicz2015empirical we add 1.0 to b_f after initialization in order to reduce the scale of forgetting in the beginning of the training. </dd> </dl> </div> Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> x = jax.random.normal(jax.random.key(0), (3, 5, 5)) >>> layer = nn.ConvLSTMCell(features=4, kernel_size=(2, 2)) >>> carry = layer.initialize_carry(jax.random.key(1), x.shape) >>> variables = layer.init(jax.random.key(2), carry, x) >>> new_carry, out = layer.apply(variables, carry, x) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell.features"> features<a class="headerlink" href="#flax.linen.ConvLSTMCell.features" title="Permalink to this definition">#</a></dt> <dd>number of convolution filters. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell.kernel_size"> kernel_size<a class="headerlink" href="#flax.linen.ConvLSTMCell.kernel_size" title="Permalink to this definition">#</a></dt> <dd>shape of the convolutional kernel. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Sequence[int] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell.strides"> strides<a class="headerlink" href="#flax.linen.ConvLSTMCell.strides" title="Permalink to this definition">#</a></dt> <dd>a sequence of <code class="docutils literal notranslate">n</code> integers, representing the inter-window strides. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Sequence[int] | None </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell.padding"> padding<a class="headerlink" href="#flax.linen.ConvLSTMCell.padding" title="Permalink to this definition">#</a></dt> <dd>either the string <code class="docutils literal notranslate">'SAME'</code>, the string <code class="docutils literal notranslate">'VALID'</code>, or a sequence of <code class="docutils literal notranslate">n</code> <code class="docutils literal notranslate">(low, high)</code> integer pairs that give the padding to apply before and after each spatial dimension. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">str | collections.abc.Sequence[tuple[int, int]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell.bias"> bias<a class="headerlink" href="#flax.linen.ConvLSTMCell.bias" title="Permalink to this definition">#</a></dt> <dd>whether to add a bias to the output (default: True). </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell.dtype"> dtype<a class="headerlink" href="#flax.linen.ConvLSTMCell.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: None). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.ConvLSTMCell.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell.__call__"> __call__(carry, inputs)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#ConvLSTMCell.__call__">[source]</a><a class="headerlink" href="#flax.linen.ConvLSTMCell.__call__" title="Permalink to this definition">#</a></dt> <dd>Constructs a convolutional LSTM. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>carry – the hidden state of the Conv2DLSTM cell, initialized using <code class="docutils literal notranslate">Conv2DLSTM.initialize_carry</code>.</li> <li>inputs – input data with dimensions (batch, spatial_dims…, features).</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A tuple with the new carry and the output. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.ConvLSTMCell.initialize_carry"> initialize_carry(rng, input_shape)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#ConvLSTMCell.initialize_carry">[source]</a><a class="headerlink" href="#flax.linen.ConvLSTMCell.initialize_carry" title="Permalink to this definition">#</a></dt> <dd>Initialize the RNN cell carry. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rng – random number generator passed to the init_fn.</li> <li>input_shape – a tuple providing the shape of the input to the cell.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">An initialized carry for the given RNN cell. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> <tr class="row-odd"><td><a class="reference internal" href="#flax.linen.ConvLSTMCell.initialize_carry" title="flax.linen.ConvLSTMCell.initialize_carry"><code class="xref py py-obj docutils literal notranslate">initialize_carry</code></a>(rng, input_shape)</td> <td>Initialize the RNN cell carry.</td> </tr> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.SimpleCell"> class flax.linen.SimpleCell(features, activation_fn=<PjitFunction of <function tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, residual=False, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#SimpleCell">[source]</a><a class="headerlink" href="#flax.linen.SimpleCell" title="Permalink to this definition">#</a></dt> <dd>Simple cell. The mathematical definition of the cell is as follows <div class="math notranslate nohighlight"> \[\begin{array}{ll} h' = \tanh(W_i x + b_i + W_h h) \end{array}\]</div> where x is the input and h is the output of the previous time step. If <cite>residual</cite> is <cite>True</cite>, <div class="math notranslate nohighlight"> \[\begin{array}{ll} h' = \tanh(W_i x + b_i + W_h h + h) \end{array}\]</div> Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> x = jax.random.normal(jax.random.key(0), (2, 3)) >>> layer = nn.SimpleCell(features=4) >>> carry = layer.initialize_carry(jax.random.key(1), x.shape) >>> variables = layer.init(jax.random.key(2), carry, x) >>> new_carry, out = layer.apply(variables, carry, x) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.features"> features<a class="headerlink" href="#flax.linen.SimpleCell.features" title="Permalink to this definition">#</a></dt> <dd>number of output features. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.activation_fn"> activation_fn<a class="headerlink" href="#flax.linen.SimpleCell.activation_fn" title="Permalink to this definition">#</a></dt> <dd>activation function used for output and memory update (default: tanh). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.SimpleCell.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the input (default: lecun_normal). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.recurrent_kernel_init"> recurrent_kernel_init<a class="headerlink" href="#flax.linen.SimpleCell.recurrent_kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.bias_init"> bias_init<a class="headerlink" href="#flax.linen.SimpleCell.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias parameters (default: initializers.zeros_init()) <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.dtype"> dtype<a class="headerlink" href="#flax.linen.SimpleCell.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: None). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.SimpleCell.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.residual"> residual<a class="headerlink" href="#flax.linen.SimpleCell.residual" title="Permalink to this definition">#</a></dt> <dd>pre-activation residual connection (<a class="reference external" href="https://arxiv.org/abs/1801.06105">https://arxiv.org/abs/1801.06105</a>). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.__call__"> __call__(carry, inputs)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#SimpleCell.__call__">[source]</a><a class="headerlink" href="#flax.linen.SimpleCell.__call__" title="Permalink to this definition">#</a></dt> <dd>Simple cell. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>carry – the hidden state of the Simple cell, initialized using <code class="docutils literal notranslate">SimpleCell.initialize_carry</code>.</li> <li>inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A tuple with the new carry and the output. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.SimpleCell.initialize_carry"> initialize_carry(rng, input_shape)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#SimpleCell.initialize_carry">[source]</a><a class="headerlink" href="#flax.linen.SimpleCell.initialize_carry" title="Permalink to this definition">#</a></dt> <dd>Initialize the RNN cell carry. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rng – random number generator passed to the init_fn.</li> <li>input_shape – a tuple providing the shape of the input to the cell.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">An initialized carry for the given RNN cell. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> <tr class="row-odd"><td><a class="reference internal" href="#flax.linen.SimpleCell.initialize_carry" title="flax.linen.SimpleCell.initialize_carry"><code class="xref py py-obj docutils literal notranslate">initialize_carry</code></a>(rng, input_shape)</td> <td>Initialize the RNN cell carry.</td> </tr> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.GRUCell"> class flax.linen.GRUCell(features, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#GRUCell">[source]</a><a class="headerlink" href="#flax.linen.GRUCell" title="Permalink to this definition">#</a></dt> <dd>GRU cell. The mathematical definition of the cell is as follows <div class="math notranslate nohighlight"> \[\begin{split}\begin{array}{ll} r = \sigma(W_{ir} x + b_{ir} + W_{hr} h) \\ z = \sigma(W_{iz} x + b_{iz} + W_{hz} h) \\ n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\ h' = (1 - z) * n + z * h \\ \end{array}\end{split}\]</div> where x is the input and h is the output of the previous time step. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> x = jax.random.normal(jax.random.key(0), (2, 3)) >>> layer = nn.GRUCell(features=4) >>> carry = layer.initialize_carry(jax.random.key(1), x.shape) >>> variables = layer.init(jax.random.key(2), carry, x) >>> new_carry, out = layer.apply(variables, carry, x) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GRUCell.features"> features<a class="headerlink" href="#flax.linen.GRUCell.features" title="Permalink to this definition">#</a></dt> <dd>number of output features. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GRUCell.gate_fn"> gate_fn<a class="headerlink" href="#flax.linen.GRUCell.gate_fn" title="Permalink to this definition">#</a></dt> <dd>activation function used for gates (default: sigmoid). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GRUCell.activation_fn"> activation_fn<a class="headerlink" href="#flax.linen.GRUCell.activation_fn" title="Permalink to this definition">#</a></dt> <dd>activation function used for output and memory update (default: tanh). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GRUCell.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.GRUCell.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the input (default: lecun_normal). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GRUCell.recurrent_kernel_init"> recurrent_kernel_init<a class="headerlink" href="#flax.linen.GRUCell.recurrent_kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GRUCell.bias_init"> bias_init<a class="headerlink" href="#flax.linen.GRUCell.bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias parameters (default: initializers.zeros_init()) <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GRUCell.dtype"> dtype<a class="headerlink" href="#flax.linen.GRUCell.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: None). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.GRUCell.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.GRUCell.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.GRUCell.__call__"> __call__(carry, inputs)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#GRUCell.__call__">[source]</a><a class="headerlink" href="#flax.linen.GRUCell.__call__" title="Permalink to this definition">#</a></dt> <dd>Gated recurrent unit (GRU) cell. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>carry – the hidden state of the GRU cell, initialized using <code class="docutils literal notranslate">GRUCell.initialize_carry</code>.</li> <li>inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A tuple with the new carry and the output. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.GRUCell.initialize_carry"> initialize_carry(rng, input_shape)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#GRUCell.initialize_carry">[source]</a><a class="headerlink" href="#flax.linen.GRUCell.initialize_carry" title="Permalink to this definition">#</a></dt> <dd>Initialize the RNN cell carry. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rng – random number generator passed to the init_fn.</li> <li>input_shape – a tuple providing the shape of the input to the cell.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">An initialized carry for the given RNN cell. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> <tr class="row-odd"><td><a class="reference internal" href="#flax.linen.GRUCell.initialize_carry" title="flax.linen.GRUCell.initialize_carry"><code class="xref py py-obj docutils literal notranslate">initialize_carry</code></a>(rng, input_shape)</td> <td>Initialize the RNN cell carry.</td> </tr> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.MGUCell"> class flax.linen.MGUCell(features, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, forget_bias_init=<function ones>, activation_bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, reset_gate=True, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#MGUCell">[source]</a><a class="headerlink" href="#flax.linen.MGUCell" title="Permalink to this definition">#</a></dt> <dd>MGU cell (<a class="reference external" href="https://arxiv.org/pdf/1603.09420.pdf">https://arxiv.org/pdf/1603.09420.pdf</a>). The mathematical definition of the cell is as follows <div class="math notranslate nohighlight"> \[\begin{split}\begin{array}{ll} f = \sigma(W_{if} x + b_{if} + W_{hf} h) \\ n = \tanh(W_{in} x + b_{in} + f * (W_{hn} h + b_{hn})) \\ h' = (1 - f) * n + f * h \\ \end{array}\end{split}\]</div> where x is the input and h is the output of the previous time step. If <code class="docutils literal notranslate">reset_gate</code> is false, the above becomes <div class="math notranslate nohighlight"> \[\begin{split}\begin{array}{ll} f = \sigma(W_{if} x + b_{if} + W_{hf} h) \\ n = \tanh(W_{in} x + b_{in} + W_{hn} h) \\ h' = (1 - f) * n + f * h \\ \end{array}\end{split}\]</div> Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> x = jax.random.normal(jax.random.key(0), (2, 3)) >>> layer = nn.MGUCell(features=4) >>> carry = layer.initialize_carry(jax.random.key(1), x.shape) >>> variables = layer.init(jax.random.key(2), carry, x) >>> new_carry, out = layer.apply(variables, carry, x) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.features"> features<a class="headerlink" href="#flax.linen.MGUCell.features" title="Permalink to this definition">#</a></dt> <dd>number of output features. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.gate_fn"> gate_fn<a class="headerlink" href="#flax.linen.MGUCell.gate_fn" title="Permalink to this definition">#</a></dt> <dd>activation function used for gates (default: sigmoid). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.activation_fn"> activation_fn<a class="headerlink" href="#flax.linen.MGUCell.activation_fn" title="Permalink to this definition">#</a></dt> <dd>activation function used for output and memory update (default: tanh). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Callable[[…], Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.kernel_init"> kernel_init<a class="headerlink" href="#flax.linen.MGUCell.kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the input (default: lecun_normal). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.recurrent_kernel_init"> recurrent_kernel_init<a class="headerlink" href="#flax.linen.MGUCell.recurrent_kernel_init" title="Permalink to this definition">#</a></dt> <dd>initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.forget_bias_init"> forget_bias_init<a class="headerlink" href="#flax.linen.MGUCell.forget_bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias parameters of the forget gate. The default is set to initializers.ones_init() because this prevents vanishing gradients. See <a class="reference external" href="https://proceedings.mlr.press/v37/jozefowicz15.pdf">https://proceedings.mlr.press/v37/jozefowicz15.pdf</a>, section 2.2 for more details. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.activation_bias_init"> activation_bias_init<a class="headerlink" href="#flax.linen.MGUCell.activation_bias_init" title="Permalink to this definition">#</a></dt> <dd>initializer for the bias parameters of the activation output (default: initializers.zeros_init()). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[jax.nn.initializers.Initializer, collections.abc.Callable[[…], Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.dtype"> dtype<a class="headerlink" href="#flax.linen.MGUCell.dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype of the computation (default: None). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.param_dtype"> param_dtype<a class="headerlink" href="#flax.linen.MGUCell.param_dtype" title="Permalink to this definition">#</a></dt> <dd>the dtype passed to parameter initializers (default: float32). <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.MGUCell.reset_gate"> reset_gate<a class="headerlink" href="#flax.linen.MGUCell.reset_gate" title="Permalink to this definition">#</a></dt> <dd>flag for applying reset gating. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.MGUCell.__call__"> __call__(carry, inputs)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#MGUCell.__call__">[source]</a><a class="headerlink" href="#flax.linen.MGUCell.__call__" title="Permalink to this definition">#</a></dt> <dd>Minimal gated unit (MGU) cell. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>carry – the hidden state of the MGU cell, initialized using <code class="docutils literal notranslate">MGUCell.initialize_carry</code>.</li> <li>inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">A tuple with the new carry and the output. </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.MGUCell.initialize_carry"> initialize_carry(rng, input_shape)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#MGUCell.initialize_carry">[source]</a><a class="headerlink" href="#flax.linen.MGUCell.initialize_carry" title="Permalink to this definition">#</a></dt> <dd>Initialize the RNN cell carry. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>rng – random number generator passed to the init_fn.</li> <li>input_shape – a tuple providing the shape of the input to the cell.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">An initialized carry for the given RNN cell. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> <tr class="row-odd"><td><a class="reference internal" href="#flax.linen.MGUCell.initialize_carry" title="flax.linen.MGUCell.initialize_carry"><code class="xref py py-obj docutils literal notranslate">initialize_carry</code></a>(rng, input_shape)</td> <td>Initialize the RNN cell carry.</td> </tr> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.RNN"> class flax.linen.RNN(cell, time_major=False, return_carry=False, reverse=False, keep_order=False, unroll=1, variable_axes=FrozenDict({}), variable_broadcast='params', variable_carry=False, split_rngs=FrozenDict({     params: False, }), parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#RNN">[source]</a><a class="headerlink" href="#flax.linen.RNN" title="Permalink to this definition">#</a></dt> <dd>The <code class="docutils literal notranslate">RNN</code> module takes any <a class="reference internal" href="#flax.linen.RNNCellBase" title="flax.linen.RNNCellBase"><code class="xref py py-class docutils literal notranslate">RNNCellBase</code></a> instance and applies it over a sequence using <a class="reference internal" href="transformations.html#flax.linen.scan" title="flax.linen.scan"><code class="xref py py-func docutils literal notranslate">flax.linen.scan()</code></a>. Example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import jax.numpy as jnp >>> import jax >>> import flax.linen as nn >>> x = jnp.ones((10, 50, 32)) # (batch, time, features) >>> lstm = nn.RNN(nn.LSTMCell(64)) >>> variables = lstm.init(jax.random.key(0), x) >>> y = lstm.apply(variables, x) >>> y.shape # (batch, time, cell_size) (10, 50, 64) </pre></div> </div> As shown above, RNN uses the <code class="docutils literal notranslate">cell_size</code> argument to set the <code class="docutils literal notranslate">size</code> argument for the cell’s <code class="docutils literal notranslate">initialize_carry</code> method, in practice this is typically the number of hidden units you want for the cell. However, this may vary depending on the cell you are using, for example the <a class="reference internal" href="#flax.linen.ConvLSTMCell" title="flax.linen.ConvLSTMCell"><code class="xref py py-class docutils literal notranslate">ConvLSTMCell</code></a> requires a <code class="docutils literal notranslate">size</code> argument of the form <code class="docutils literal notranslate">(kernel_height, kernel_width, features)</code>: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> x = jnp.ones((10, 50, 32, 32, 3)) # (batch, time, height, width, features) >>> conv_lstm = nn.RNN(nn.ConvLSTMCell(64, kernel_size=(3, 3))) >>> y, variables = conv_lstm.init_with_output(jax.random.key(0), x) >>> y.shape # (batch, time, height, width, features) (10, 50, 32, 32, 64) </pre></div> </div> By default RNN expect the time dimension after the batch dimension (<code class="docutils literal notranslate">(*batch, time, *features)</code>), if you set <code class="docutils literal notranslate">time_major=True</code> RNN will instead expect the time dimesion to be at the beginning (<code class="docutils literal notranslate">(time, *batch, *features)</code>): <div class="highlight-default notranslate"><div class="highlight"><pre>>>> x = jnp.ones((50, 10, 32)) # (time, batch, features) >>> lstm = nn.RNN(nn.LSTMCell(64), time_major=True) >>> variables = lstm.init(jax.random.key(0), x) >>> y = lstm.apply(variables, x) >>> y.shape # (time, batch, cell_size) (50, 10, 64) </pre></div> </div> The output is an array of shape <code class="docutils literal notranslate">(*batch, time, *cell_size)</code> by default (typically), however if you set <code class="docutils literal notranslate">return_carry=True</code> it will instead return a tuple of the final carry and the output: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> x = jnp.ones((10, 50, 32)) # (batch, time, features) >>> lstm = nn.RNN(nn.LSTMCell(64), return_carry=True) >>> variables = lstm.init(jax.random.key(0), x) >>> carry, y = lstm.apply(variables, x) >>> jax.tree_util.tree_map(jnp.shape, carry) # ((batch, cell_size), (batch, cell_size)) ((10, 64), (10, 64)) >>> y.shape # (batch, time, cell_size) (10, 50, 64) </pre></div> </div> To support variable length sequences, you can pass a <code class="docutils literal notranslate">seq_lengths</code> which is an integer array of shape <code class="docutils literal notranslate">(*batch)</code> where each element is the length of the sequence in the batch. For example: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> seq_lengths = jnp.array([3, 2, 5]) </pre></div> </div> The output elements corresponding to padding elements are NOT zeroed out. If <code class="docutils literal notranslate">return_carry</code> is set to <code class="docutils literal notranslate">True</code> the carry will be the state of the last valid element of each sequence. RNN also accepts some of the arguments of <a class="reference internal" href="transformations.html#flax.linen.scan" title="flax.linen.scan"><code class="xref py py-func docutils literal notranslate">flax.linen.scan()</code></a>, by default they are set to work with cells like <a class="reference internal" href="#flax.linen.LSTMCell" title="flax.linen.LSTMCell"><code class="xref py py-class docutils literal notranslate">LSTMCell</code></a> and <a class="reference internal" href="#flax.linen.GRUCell" title="flax.linen.GRUCell"><code class="xref py py-class docutils literal notranslate">GRUCell</code></a> but they can be overriden as needed. Overriding default values to scan looks like this: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> lstm = nn.RNN( ... nn.LSTMCell(64), ... unroll=1, variable_axes={}, variable_broadcast='params', ... variable_carry=False, split_rngs={'params': False}) </pre></div> </div> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.cell"> cell<a class="headerlink" href="#flax.linen.RNN.cell" title="Permalink to this definition">#</a></dt> <dd>an instance of <a class="reference internal" href="#flax.linen.RNNCellBase" title="flax.linen.RNNCellBase"><code class="xref py py-class docutils literal notranslate">RNNCellBase</code></a>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd"><a class="reference internal" href="#flax.linen.RNNCellBase" title="flax.linen.recurrent.RNNCellBase">flax.linen.recurrent.RNNCellBase</a> </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.time_major"> time_major<a class="headerlink" href="#flax.linen.RNN.time_major" title="Permalink to this definition">#</a></dt> <dd>if <code class="docutils literal notranslate">time_major=False</code> (default) it will expect inputs with shape <code class="docutils literal notranslate">(*batch, time, *features)</code>, else it will expect inputs with shape <code class="docutils literal notranslate">(time, *batch, *features)</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.return_carry"> return_carry<a class="headerlink" href="#flax.linen.RNN.return_carry" title="Permalink to this definition">#</a></dt> <dd>if <code class="docutils literal notranslate">return_carry=False</code> (default) only the output sequence is returned, else it will return a tuple of the final carry and the output sequence. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.reverse"> reverse<a class="headerlink" href="#flax.linen.RNN.reverse" title="Permalink to this definition">#</a></dt> <dd>if <code class="docutils literal notranslate">reverse=False</code> (default) the sequence is processed from left to right and returned in the original order, else it will be processed from right to left, and returned in reverse order. If <code class="docutils literal notranslate">seq_lengths</code> is passed, padding will always remain at the end of the sequence. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.keep_order"> keep_order<a class="headerlink" href="#flax.linen.RNN.keep_order" title="Permalink to this definition">#</a></dt> <dd>if <code class="docutils literal notranslate">keep_order=True</code>, when <code class="docutils literal notranslate">reverse=True</code> the output will be reversed back to the original order after processing, this is useful to align sequences in bidirectional RNNs. If <code class="docutils literal notranslate">keep_order=False</code> (default), the output will remain in the order specified by <code class="docutils literal notranslate">reverse</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">bool </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.unroll"> unroll<a class="headerlink" href="#flax.linen.RNN.unroll" title="Permalink to this definition">#</a></dt> <dd>how many scan iterations to unroll within a single iteration of a loop, defaults to 1. This argument will be passed to <code class="docutils literal notranslate">nn.scan</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">int </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.variable_axes"> variable_axes<a class="headerlink" href="#flax.linen.RNN.variable_axes" title="Permalink to this definition">#</a></dt> <dd>a dictionary mapping each collection to either an integer <code class="docutils literal notranslate">i</code> (meaning we scan over dimension <code class="docutils literal notranslate">i</code>) or <code class="docutils literal notranslate">None</code> (replicate rather than scan). This argument is forwarded to <code class="docutils literal notranslate">nn.scan</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Mapping[Union[bool, str, Collection[str], DenyList], Union[int, flax.typing.In[int], flax.typing.Out[int]]] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.variable_broadcast"> variable_broadcast<a class="headerlink" href="#flax.linen.RNN.variable_broadcast" title="Permalink to this definition">#</a></dt> <dd>Specifies the broadcasted variable collections. A broadcasted variable should not depend on any computation that cannot be lifted out of the loop. This is typically used to define shared parameters inside the fn. This argument is forwarded to <code class="docutils literal notranslate">nn.scan</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[bool, str, Collection[str], DenyList] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.variable_carry"> variable_carry<a class="headerlink" href="#flax.linen.RNN.variable_carry" title="Permalink to this definition">#</a></dt> <dd>Specifies the variable collections that are carried through the loop. Mutations to these variables are carried to the next iteration and will be preserved when the scan finishes. This argument is forwarded to <code class="docutils literal notranslate">nn.scan</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">Union[bool, str, Collection[str], DenyList] </dd> </dl> </dd></dl> <dl class="py attribute"> <dt class="sig sig-object py" id="flax.linen.RNN.split_rngs"> split_rngs<a class="headerlink" href="#flax.linen.RNN.split_rngs" title="Permalink to this definition">#</a></dt> <dd>a mapping from PRNGSequenceFilter to bool specifying whether a collection’s PRNG key should be split such that its values are different at each step, or replicated such that its values remain the same at each step. This argument is forwarded to <code class="docutils literal notranslate">nn.scan</code>. <dl class="field-list simple"> <dt class="field-odd">Type</dt> <dd class="field-odd">collections.abc.Mapping[Union[bool, str, Collection[str], DenyList], bool] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.RNN.__call__"> __call__(inputs, *, initial_carry=None, init_key=None, seq_lengths=None, return_carry=None, time_major=None, reverse=None, keep_order=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#RNN.__call__">[source]</a><a class="headerlink" href="#flax.linen.RNN.__call__" title="Permalink to this definition">#</a></dt> <dd>Applies the RNN to the inputs. <code class="docutils literal notranslate">__call__</code> allows you to optionally override some attributes like <code class="docutils literal notranslate">return_carry</code> and <code class="docutils literal notranslate">time_major</code> defined in the constructor. <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li>inputs – the input sequence.</li> <li>initial_carry – the initial carry, if not provided it will be initialized using the cell’s <a class="reference internal" href="#flax.linen.RNNCellBase.initialize_carry" title="flax.linen.RNNCellBase.initialize_carry"><code class="xref py py-meth docutils literal notranslate">RNNCellBase.initialize_carry()</code></a> method.</li> <li>init_key – a PRNG key used to initialize the carry, if not provided <code class="docutils literal notranslate">jax.random.key(0)</code> will be used. Most cells will ignore this argument.</li> <li>seq_lengths – an optional integer array of shape <code class="docutils literal notranslate">(*batch)</code> indicating the length of each sequence, elements whose index in the time dimension is greater than the corresponding length will be considered padding and will be ignored.</li> <li>return_carry – if <code class="docutils literal notranslate">return_carry=False</code> (default) only the output sequence is returned, else it will return a tuple of the final carry and the output sequence.</li> <li>time_major – if <code class="docutils literal notranslate">time_major=False</code> (default) it will expect inputs with shape <code class="docutils literal notranslate">(*batch, time, *features)</code>, else it will expect inputs with shape <code class="docutils literal notranslate">(time, *batch, *features)</code>.</li> <li>reverse – overrides the <code class="docutils literal notranslate">reverse</code> attribute, if <code class="docutils literal notranslate">reverse=False</code> (default) the sequence is processed from left to right and returned in the original order, else it will be processed from right to left, and returned in reverse order. If <code class="docutils literal notranslate">seq_lengths</code> is passed, padding will always remain at the end of the sequence.</li> <li>keep_order – overrides the <code class="docutils literal notranslate">keep_order</code> attribute, if <code class="docutils literal notranslate">keep_order=True</code>, when <code class="docutils literal notranslate">reverse=True</code> the output will be reversed back to the original order after processing, this is useful to align sequences in bidirectional RNNs. If <code class="docutils literal notranslate">keep_order=False</code> (default), the output will remain in the order specified by <code class="docutils literal notranslate">reverse</code>.</li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even">if <code class="docutils literal notranslate">return_carry=False</code> (default) only the output sequence is returned, else it will return a tuple of the final carry and the output sequence. </dd> </dl> </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.Bidirectional"> class flax.linen.Bidirectional(forward_rnn, backward_rnn, merge_fn=<function _concatenate>, time_major=False, return_carry=False, parent=<flax.linen.module._Sentinel object>, name=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#Bidirectional">[source]</a><a class="headerlink" href="#flax.linen.Bidirectional" title="Permalink to this definition">#</a></dt> <dd>Processes the input in both directions and merges the results. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import flax.linen as nn >>> import jax, jax.numpy as jnp >>> layer = nn.Bidirectional(nn.RNN(nn.GRUCell(4)), nn.RNN(nn.GRUCell(4))) >>> x = jnp.ones((2, 3)) >>> variables = layer.init(jax.random.key(0), x) >>> out = layer.apply(variables, x) </pre></div> </div> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.Bidirectional.__call__"> __call__(inputs, *, initial_carry=None, init_key=None, seq_lengths=None, return_carry=None, time_major=None, reverse=None, keep_order=None)<a class="reference internal" href="../../_modules/flax/linen/recurrent.html#Bidirectional.__call__">[source]</a><a class="headerlink" href="#flax.linen.Bidirectional.__call__" title="Permalink to this definition">#</a></dt> <dd>Call self as a function. </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> </div> <div class="section" id="batchapply"> <h2>BatchApply<a class="headerlink" href="#batchapply" title="Permalink to this heading">#</a></h2> <div class="docutils container"> <dl class="py class"> <dt class="sig sig-object py" id="flax.linen.BatchApply"> class flax.linen.BatchApply(f, num_dims=2)<a class="reference internal" href="../../_modules/flax/linen/batch_apply.html#BatchApply">[source]</a><a class="headerlink" href="#flax.linen.BatchApply" title="Permalink to this definition">#</a></dt> <dd>Temporarily merges leading dimensions of input tensors. Merges the leading dimensions of a tensor into a single dimension, runs the given callable, then splits the leading dimension of the result to match the input. Input arrays whose rank is smaller than the number of dimensions to collapse are passed unmodified. This may be useful for applying a module to each timestep of e.g. a <code class="docutils literal notranslate">[Time, Batch, ...]</code> array. For some <code class="docutils literal notranslate">f</code>s and platforms, this may be more efficient than <code class="xref py py-func docutils literal notranslate">jax.vmap()</code>, especially when combined with other transformations like <code class="xref py py-func docutils literal notranslate">jax.grad()</code>. Example usage: <div class="highlight-default notranslate"><div class="highlight"><pre>>>> import jax, jax.numpy as jnp >>> a = jax.random.normal(jax.random.key(0), [2, 3, 4]) >>> b = jax.random.normal(jax.random.key(1), [4]) >>> def raises(a, b): ... if len(a.shape) != 2: ... raise ValueError("a must be shape 2") ... if len(b.shape) != 1: ... raise ValueError("b must be shape 1") ... return jnp.dot(a, b) >>> out = BatchApply(raises)(a, b) >>> expected_merged_leading = raises(a.reshape(2*3, 4), b) >>> expected = expected_merged_leading.reshape((2, 3) + expected_merged_leading.shape[1:]) >>> np.testing.assert_array_equal(out, expected) </pre></div> </div> <dl class="py method"> <dt class="sig sig-object py" id="flax.linen.BatchApply.__call__"> __call__(*args, **kwargs)<a class="reference internal" href="../../_modules/flax/linen/batch_apply.html#BatchApply.__call__">[source]</a><a class="headerlink" href="#flax.linen.BatchApply.__call__" title="Permalink to this definition">#</a></dt> <dd>Call self as a function. </dd></dl> Methods <div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary"> <colgroup> <col style="width: 10%" /> <col style="width: 90%" /> </colgroup> <tbody> </tbody> </table> </div> </dd></dl> </div> </div> </div> </article> <footer class="prev-next-footer d-print-none"> <div class="prev-next-area"> <a class="left-prev" href="init_apply.html" title="previous page"> <div class="prev-next-info"> previous Init/Apply </div> </a> <a class="right-next" href="activation_functions.html" title="next page"> <div class="prev-next-info"> next Activation functions </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="#linear-modules">Linear Modules</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.Dense"><code class="docutils literal notranslate">Dense</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.features"><code class="docutils literal notranslate">Dense.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.use_bias"><code class="docutils literal notranslate">Dense.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.dtype"><code class="docutils literal notranslate">Dense.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.param_dtype"><code class="docutils literal notranslate">Dense.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.precision"><code class="docutils literal notranslate">Dense.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.kernel_init"><code class="docutils literal notranslate">Dense.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.bias_init"><code class="docutils literal notranslate">Dense.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dense.__call__"><code class="docutils literal notranslate">Dense.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral"><code class="docutils literal notranslate">DenseGeneral</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.features"><code class="docutils literal notranslate">DenseGeneral.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.axis"><code class="docutils literal notranslate">DenseGeneral.axis</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.batch_dims"><code class="docutils literal notranslate">DenseGeneral.batch_dims</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.use_bias"><code class="docutils literal notranslate">DenseGeneral.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.dtype"><code class="docutils literal notranslate">DenseGeneral.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.param_dtype"><code class="docutils literal notranslate">DenseGeneral.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.kernel_init"><code class="docutils literal notranslate">DenseGeneral.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.bias_init"><code class="docutils literal notranslate">DenseGeneral.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.precision"><code class="docutils literal notranslate">DenseGeneral.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.DenseGeneral.__call__"><code class="docutils literal notranslate">DenseGeneral.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv"><code class="docutils literal notranslate">Conv</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.features"><code class="docutils literal notranslate">Conv.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.kernel_size"><code class="docutils literal notranslate">Conv.kernel_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.strides"><code class="docutils literal notranslate">Conv.strides</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.padding"><code class="docutils literal notranslate">Conv.padding</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.input_dilation"><code class="docutils literal notranslate">Conv.input_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.kernel_dilation"><code class="docutils literal notranslate">Conv.kernel_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.feature_group_count"><code class="docutils literal notranslate">Conv.feature_group_count</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.use_bias"><code class="docutils literal notranslate">Conv.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.mask"><code class="docutils literal notranslate">Conv.mask</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.dtype"><code class="docutils literal notranslate">Conv.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.param_dtype"><code class="docutils literal notranslate">Conv.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.precision"><code class="docutils literal notranslate">Conv.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.kernel_init"><code class="docutils literal notranslate">Conv.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.bias_init"><code class="docutils literal notranslate">Conv.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Conv.__call__"><code class="docutils literal notranslate">Conv.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose"><code class="docutils literal notranslate">ConvTranspose</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.features"><code class="docutils literal notranslate">ConvTranspose.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.kernel_size"><code class="docutils literal notranslate">ConvTranspose.kernel_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.strides"><code class="docutils literal notranslate">ConvTranspose.strides</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.padding"><code class="docutils literal notranslate">ConvTranspose.padding</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.kernel_dilation"><code class="docutils literal notranslate">ConvTranspose.kernel_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.use_bias"><code class="docutils literal notranslate">ConvTranspose.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.mask"><code class="docutils literal notranslate">ConvTranspose.mask</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.dtype"><code class="docutils literal notranslate">ConvTranspose.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.param_dtype"><code class="docutils literal notranslate">ConvTranspose.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.precision"><code class="docutils literal notranslate">ConvTranspose.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.kernel_init"><code class="docutils literal notranslate">ConvTranspose.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.bias_init"><code class="docutils literal notranslate">ConvTranspose.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.transpose_kernel"><code class="docutils literal notranslate">ConvTranspose.transpose_kernel</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvTranspose.__call__"><code class="docutils literal notranslate">ConvTranspose.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal"><code class="docutils literal notranslate">ConvLocal</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.features"><code class="docutils literal notranslate">ConvLocal.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.kernel_size"><code class="docutils literal notranslate">ConvLocal.kernel_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.strides"><code class="docutils literal notranslate">ConvLocal.strides</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.padding"><code class="docutils literal notranslate">ConvLocal.padding</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.input_dilation"><code class="docutils literal notranslate">ConvLocal.input_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.kernel_dilation"><code class="docutils literal notranslate">ConvLocal.kernel_dilation</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.feature_group_count"><code class="docutils literal notranslate">ConvLocal.feature_group_count</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.use_bias"><code class="docutils literal notranslate">ConvLocal.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.mask"><code class="docutils literal notranslate">ConvLocal.mask</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.dtype"><code class="docutils literal notranslate">ConvLocal.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.param_dtype"><code class="docutils literal notranslate">ConvLocal.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.precision"><code class="docutils literal notranslate">ConvLocal.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.kernel_init"><code class="docutils literal notranslate">ConvLocal.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.bias_init"><code class="docutils literal notranslate">ConvLocal.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLocal.__call__"><code class="docutils literal notranslate">ConvLocal.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum"><code class="docutils literal notranslate">Einsum</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.shape"><code class="docutils literal notranslate">Einsum.shape</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.einsum_str"><code class="docutils literal notranslate">Einsum.einsum_str</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.use_bias"><code class="docutils literal notranslate">Einsum.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.dtype"><code class="docutils literal notranslate">Einsum.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.param_dtype"><code class="docutils literal notranslate">Einsum.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.precision"><code class="docutils literal notranslate">Einsum.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.kernel_init"><code class="docutils literal notranslate">Einsum.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.bias_init"><code class="docutils literal notranslate">Einsum.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Einsum.__call__"><code class="docutils literal notranslate">Einsum.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed"><code class="docutils literal notranslate">Embed</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.num_embeddings"><code class="docutils literal notranslate">Embed.num_embeddings</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.features"><code class="docutils literal notranslate">Embed.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.dtype"><code class="docutils literal notranslate">Embed.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.param_dtype"><code class="docutils literal notranslate">Embed.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.embedding_init"><code class="docutils literal notranslate">Embed.embedding_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.__call__"><code class="docutils literal notranslate">Embed.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Embed.attend"><code class="docutils literal notranslate">Embed.attend()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#pooling">Pooling</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.max_pool"><code class="docutils literal notranslate">max_pool()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.avg_pool"><code class="docutils literal notranslate">avg_pool()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.pool"><code class="docutils literal notranslate">pool()</code></a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#normalization">Normalization</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.BatchNorm"><code class="docutils literal notranslate">BatchNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.use_running_average"><code class="docutils literal notranslate">BatchNorm.use_running_average</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.axis"><code class="docutils literal notranslate">BatchNorm.axis</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.momentum"><code class="docutils literal notranslate">BatchNorm.momentum</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.epsilon"><code class="docutils literal notranslate">BatchNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.dtype"><code class="docutils literal notranslate">BatchNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.param_dtype"><code class="docutils literal notranslate">BatchNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.use_bias"><code class="docutils literal notranslate">BatchNorm.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.use_scale"><code class="docutils literal notranslate">BatchNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.bias_init"><code class="docutils literal notranslate">BatchNorm.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.scale_init"><code class="docutils literal notranslate">BatchNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.axis_name"><code class="docutils literal notranslate">BatchNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.axis_index_groups"><code class="docutils literal notranslate">BatchNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.use_fast_variance"><code class="docutils literal notranslate">BatchNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchNorm.__call__"><code class="docutils literal notranslate">BatchNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm"><code class="docutils literal notranslate">LayerNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.epsilon"><code class="docutils literal notranslate">LayerNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.dtype"><code class="docutils literal notranslate">LayerNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.param_dtype"><code class="docutils literal notranslate">LayerNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.use_bias"><code class="docutils literal notranslate">LayerNorm.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.use_scale"><code class="docutils literal notranslate">LayerNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.bias_init"><code class="docutils literal notranslate">LayerNorm.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.scale_init"><code class="docutils literal notranslate">LayerNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.reduction_axes"><code class="docutils literal notranslate">LayerNorm.reduction_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.feature_axes"><code class="docutils literal notranslate">LayerNorm.feature_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.axis_name"><code class="docutils literal notranslate">LayerNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.axis_index_groups"><code class="docutils literal notranslate">LayerNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.use_fast_variance"><code class="docutils literal notranslate">LayerNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LayerNorm.__call__"><code class="docutils literal notranslate">LayerNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm"><code class="docutils literal notranslate">GroupNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.num_groups"><code class="docutils literal notranslate">GroupNorm.num_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.group_size"><code class="docutils literal notranslate">GroupNorm.group_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.epsilon"><code class="docutils literal notranslate">GroupNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.dtype"><code class="docutils literal notranslate">GroupNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.param_dtype"><code class="docutils literal notranslate">GroupNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.use_bias"><code class="docutils literal notranslate">GroupNorm.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.use_scale"><code class="docutils literal notranslate">GroupNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.bias_init"><code class="docutils literal notranslate">GroupNorm.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.scale_init"><code class="docutils literal notranslate">GroupNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.reduction_axes"><code class="docutils literal notranslate">GroupNorm.reduction_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.axis_name"><code class="docutils literal notranslate">GroupNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.axis_index_groups"><code class="docutils literal notranslate">GroupNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.use_fast_variance"><code class="docutils literal notranslate">GroupNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GroupNorm.__call__"><code class="docutils literal notranslate">GroupNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm"><code class="docutils literal notranslate">RMSNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.epsilon"><code class="docutils literal notranslate">RMSNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.dtype"><code class="docutils literal notranslate">RMSNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.param_dtype"><code class="docutils literal notranslate">RMSNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.use_scale"><code class="docutils literal notranslate">RMSNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.scale_init"><code class="docutils literal notranslate">RMSNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.reduction_axes"><code class="docutils literal notranslate">RMSNorm.reduction_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.feature_axes"><code class="docutils literal notranslate">RMSNorm.feature_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.axis_name"><code class="docutils literal notranslate">RMSNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.axis_index_groups"><code class="docutils literal notranslate">RMSNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.use_fast_variance"><code class="docutils literal notranslate">RMSNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RMSNorm.__call__"><code class="docutils literal notranslate">RMSNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm"><code class="docutils literal notranslate">InstanceNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.epsilon"><code class="docutils literal notranslate">InstanceNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.dtype"><code class="docutils literal notranslate">InstanceNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.param_dtype"><code class="docutils literal notranslate">InstanceNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.use_bias"><code class="docutils literal notranslate">InstanceNorm.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.use_scale"><code class="docutils literal notranslate">InstanceNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.bias_init"><code class="docutils literal notranslate">InstanceNorm.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.scale_init"><code class="docutils literal notranslate">InstanceNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.feature_axes"><code class="docutils literal notranslate">InstanceNorm.feature_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.axis_name"><code class="docutils literal notranslate">InstanceNorm.axis_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.axis_index_groups"><code class="docutils literal notranslate">InstanceNorm.axis_index_groups</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.use_fast_variance"><code class="docutils literal notranslate">InstanceNorm.use_fast_variance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.InstanceNorm.__call__"><code class="docutils literal notranslate">InstanceNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm"><code class="docutils literal notranslate">SpectralNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.layer_instance"><code class="docutils literal notranslate">SpectralNorm.layer_instance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.n_steps"><code class="docutils literal notranslate">SpectralNorm.n_steps</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.epsilon"><code class="docutils literal notranslate">SpectralNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.dtype"><code class="docutils literal notranslate">SpectralNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.param_dtype"><code class="docutils literal notranslate">SpectralNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.error_on_non_matrix"><code class="docutils literal notranslate">SpectralNorm.error_on_non_matrix</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.collection_name"><code class="docutils literal notranslate">SpectralNorm.collection_name</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SpectralNorm.__call__"><code class="docutils literal notranslate">SpectralNorm.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm"><code class="docutils literal notranslate">WeightNorm</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.layer_instance"><code class="docutils literal notranslate">WeightNorm.layer_instance</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.epsilon"><code class="docutils literal notranslate">WeightNorm.epsilon</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.dtype"><code class="docutils literal notranslate">WeightNorm.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.param_dtype"><code class="docutils literal notranslate">WeightNorm.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.use_scale"><code class="docutils literal notranslate">WeightNorm.use_scale</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.scale_init"><code class="docutils literal notranslate">WeightNorm.scale_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.feature_axes"><code class="docutils literal notranslate">WeightNorm.feature_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.variable_filter"><code class="docutils literal notranslate">WeightNorm.variable_filter</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.WeightNorm.__call__"><code class="docutils literal notranslate">WeightNorm.__call__()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#combinators">Combinators</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.Sequential"><code class="docutils literal notranslate">Sequential</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Sequential.layers"><code class="docutils literal notranslate">Sequential.layers</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Sequential.__call__"><code class="docutils literal notranslate">Sequential.__call__()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#stochastic">Stochastic</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.Dropout"><code class="docutils literal notranslate">Dropout</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.rate"><code class="docutils literal notranslate">Dropout.rate</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.broadcast_dims"><code class="docutils literal notranslate">Dropout.broadcast_dims</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.deterministic"><code class="docutils literal notranslate">Dropout.deterministic</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.rng_collection"><code class="docutils literal notranslate">Dropout.rng_collection</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Dropout.__call__"><code class="docutils literal notranslate">Dropout.__call__()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#attention">Attention</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.MultiHeadDotProductAttention"><code class="docutils literal notranslate">MultiHeadDotProductAttention</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.num_heads"><code class="docutils literal notranslate">MultiHeadDotProductAttention.num_heads</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.dtype"><code class="docutils literal notranslate">MultiHeadDotProductAttention.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.param_dtype"><code class="docutils literal notranslate">MultiHeadDotProductAttention.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.qkv_features"><code class="docutils literal notranslate">MultiHeadDotProductAttention.qkv_features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.out_features"><code class="docutils literal notranslate">MultiHeadDotProductAttention.out_features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.broadcast_dropout"><code class="docutils literal notranslate">MultiHeadDotProductAttention.broadcast_dropout</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.dropout_rate"><code class="docutils literal notranslate">MultiHeadDotProductAttention.dropout_rate</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.deterministic"><code class="docutils literal notranslate">MultiHeadDotProductAttention.deterministic</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.precision"><code class="docutils literal notranslate">MultiHeadDotProductAttention.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.kernel_init"><code class="docutils literal notranslate">MultiHeadDotProductAttention.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.out_kernel_init"><code class="docutils literal notranslate">MultiHeadDotProductAttention.out_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.bias_init"><code class="docutils literal notranslate">MultiHeadDotProductAttention.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.out_bias_init"><code class="docutils literal notranslate">MultiHeadDotProductAttention.out_bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.use_bias"><code class="docutils literal notranslate">MultiHeadDotProductAttention.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.attention_fn"><code class="docutils literal notranslate">MultiHeadDotProductAttention.attention_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.decode"><code class="docutils literal notranslate">MultiHeadDotProductAttention.decode</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.normalize_qk"><code class="docutils literal notranslate">MultiHeadDotProductAttention.normalize_qk</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.qk_attn_weights_einsum_cls"><code class="docutils literal notranslate">MultiHeadDotProductAttention.qk_attn_weights_einsum_cls</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.attn_weights_value_einsum_cls"><code class="docutils literal notranslate">MultiHeadDotProductAttention.attn_weights_value_einsum_cls</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadDotProductAttention.__call__"><code class="docutils literal notranslate">MultiHeadDotProductAttention.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention"><code class="docutils literal notranslate">MultiHeadAttention</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.num_heads"><code class="docutils literal notranslate">MultiHeadAttention.num_heads</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.dtype"><code class="docutils literal notranslate">MultiHeadAttention.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.param_dtype"><code class="docutils literal notranslate">MultiHeadAttention.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.qkv_features"><code class="docutils literal notranslate">MultiHeadAttention.qkv_features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.out_features"><code class="docutils literal notranslate">MultiHeadAttention.out_features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.broadcast_dropout"><code class="docutils literal notranslate">MultiHeadAttention.broadcast_dropout</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.dropout_rate"><code class="docutils literal notranslate">MultiHeadAttention.dropout_rate</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.deterministic"><code class="docutils literal notranslate">MultiHeadAttention.deterministic</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.precision"><code class="docutils literal notranslate">MultiHeadAttention.precision</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.kernel_init"><code class="docutils literal notranslate">MultiHeadAttention.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.bias_init"><code class="docutils literal notranslate">MultiHeadAttention.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.use_bias"><code class="docutils literal notranslate">MultiHeadAttention.use_bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.attention_fn"><code class="docutils literal notranslate">MultiHeadAttention.attention_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.decode"><code class="docutils literal notranslate">MultiHeadAttention.decode</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.normalize_qk"><code class="docutils literal notranslate">MultiHeadAttention.normalize_qk</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MultiHeadAttention.__call__"><code class="docutils literal notranslate">MultiHeadAttention.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SelfAttention"><code class="docutils literal notranslate">SelfAttention</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SelfAttention.__call__"><code class="docutils literal notranslate">SelfAttention.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.dot_product_attention_weights"><code class="docutils literal notranslate">dot_product_attention_weights()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.dot_product_attention"><code class="docutils literal notranslate">dot_product_attention()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.make_attention_mask"><code class="docutils literal notranslate">make_attention_mask()</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.make_causal_mask"><code class="docutils literal notranslate">make_causal_mask()</code></a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#recurrent">Recurrent</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.RNNCellBase"><code class="docutils literal notranslate">RNNCellBase</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNNCellBase.__call__"><code class="docutils literal notranslate">RNNCellBase.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNNCellBase.initialize_carry"><code class="docutils literal notranslate">RNNCellBase.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell"><code class="docutils literal notranslate">LSTMCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.features"><code class="docutils literal notranslate">LSTMCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.gate_fn"><code class="docutils literal notranslate">LSTMCell.gate_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.activation_fn"><code class="docutils literal notranslate">LSTMCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.kernel_init"><code class="docutils literal notranslate">LSTMCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.recurrent_kernel_init"><code class="docutils literal notranslate">LSTMCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.bias_init"><code class="docutils literal notranslate">LSTMCell.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.dtype"><code class="docutils literal notranslate">LSTMCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.param_dtype"><code class="docutils literal notranslate">LSTMCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.__call__"><code class="docutils literal notranslate">LSTMCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.LSTMCell.initialize_carry"><code class="docutils literal notranslate">LSTMCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell"><code class="docutils literal notranslate">OptimizedLSTMCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.gate_fn"><code class="docutils literal notranslate">OptimizedLSTMCell.gate_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.activation_fn"><code class="docutils literal notranslate">OptimizedLSTMCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.kernel_init"><code class="docutils literal notranslate">OptimizedLSTMCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.recurrent_kernel_init"><code class="docutils literal notranslate">OptimizedLSTMCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.bias_init"><code class="docutils literal notranslate">OptimizedLSTMCell.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.dtype"><code class="docutils literal notranslate">OptimizedLSTMCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.param_dtype"><code class="docutils literal notranslate">OptimizedLSTMCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.__call__"><code class="docutils literal notranslate">OptimizedLSTMCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.OptimizedLSTMCell.initialize_carry"><code class="docutils literal notranslate">OptimizedLSTMCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell"><code class="docutils literal notranslate">ConvLSTMCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.features"><code class="docutils literal notranslate">ConvLSTMCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.kernel_size"><code class="docutils literal notranslate">ConvLSTMCell.kernel_size</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.strides"><code class="docutils literal notranslate">ConvLSTMCell.strides</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.padding"><code class="docutils literal notranslate">ConvLSTMCell.padding</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.bias"><code class="docutils literal notranslate">ConvLSTMCell.bias</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.dtype"><code class="docutils literal notranslate">ConvLSTMCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.param_dtype"><code class="docutils literal notranslate">ConvLSTMCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.__call__"><code class="docutils literal notranslate">ConvLSTMCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.ConvLSTMCell.initialize_carry"><code class="docutils literal notranslate">ConvLSTMCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell"><code class="docutils literal notranslate">SimpleCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.features"><code class="docutils literal notranslate">SimpleCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.activation_fn"><code class="docutils literal notranslate">SimpleCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.kernel_init"><code class="docutils literal notranslate">SimpleCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.recurrent_kernel_init"><code class="docutils literal notranslate">SimpleCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.bias_init"><code class="docutils literal notranslate">SimpleCell.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.dtype"><code class="docutils literal notranslate">SimpleCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.param_dtype"><code class="docutils literal notranslate">SimpleCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.residual"><code class="docutils literal notranslate">SimpleCell.residual</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.__call__"><code class="docutils literal notranslate">SimpleCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.SimpleCell.initialize_carry"><code class="docutils literal notranslate">SimpleCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell"><code class="docutils literal notranslate">GRUCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.features"><code class="docutils literal notranslate">GRUCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.gate_fn"><code class="docutils literal notranslate">GRUCell.gate_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.activation_fn"><code class="docutils literal notranslate">GRUCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.kernel_init"><code class="docutils literal notranslate">GRUCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.recurrent_kernel_init"><code class="docutils literal notranslate">GRUCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.bias_init"><code class="docutils literal notranslate">GRUCell.bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.dtype"><code class="docutils literal notranslate">GRUCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.param_dtype"><code class="docutils literal notranslate">GRUCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.__call__"><code class="docutils literal notranslate">GRUCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.GRUCell.initialize_carry"><code class="docutils literal notranslate">GRUCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell"><code class="docutils literal notranslate">MGUCell</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.features"><code class="docutils literal notranslate">MGUCell.features</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.gate_fn"><code class="docutils literal notranslate">MGUCell.gate_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.activation_fn"><code class="docutils literal notranslate">MGUCell.activation_fn</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.kernel_init"><code class="docutils literal notranslate">MGUCell.kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.recurrent_kernel_init"><code class="docutils literal notranslate">MGUCell.recurrent_kernel_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.forget_bias_init"><code class="docutils literal notranslate">MGUCell.forget_bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.activation_bias_init"><code class="docutils literal notranslate">MGUCell.activation_bias_init</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.dtype"><code class="docutils literal notranslate">MGUCell.dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.param_dtype"><code class="docutils literal notranslate">MGUCell.param_dtype</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.reset_gate"><code class="docutils literal notranslate">MGUCell.reset_gate</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.__call__"><code class="docutils literal notranslate">MGUCell.__call__()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.MGUCell.initialize_carry"><code class="docutils literal notranslate">MGUCell.initialize_carry()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN"><code class="docutils literal notranslate">RNN</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.cell"><code class="docutils literal notranslate">RNN.cell</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.time_major"><code class="docutils literal notranslate">RNN.time_major</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.return_carry"><code class="docutils literal notranslate">RNN.return_carry</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.reverse"><code class="docutils literal notranslate">RNN.reverse</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.keep_order"><code class="docutils literal notranslate">RNN.keep_order</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.unroll"><code class="docutils literal notranslate">RNN.unroll</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.variable_axes"><code class="docutils literal notranslate">RNN.variable_axes</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.variable_broadcast"><code class="docutils literal notranslate">RNN.variable_broadcast</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.variable_carry"><code class="docutils literal notranslate">RNN.variable_carry</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.split_rngs"><code class="docutils literal notranslate">RNN.split_rngs</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.RNN.__call__"><code class="docutils literal notranslate">RNN.__call__()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Bidirectional"><code class="docutils literal notranslate">Bidirectional</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.Bidirectional.__call__"><code class="docutils literal notranslate">Bidirectional.__call__()</code></a></li> </ul> </li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#batchapply">BatchApply</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.BatchApply"><code class="docutils literal notranslate">BatchApply</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.linen.BatchApply.__call__"><code class="docutils literal notranslate">BatchApply.__call__()</code></a></li> </ul> </li> </ul> </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>