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>Migrating from Haiku to Flax</title> <script data-cfasync="false"> document.documentElement.dataset.mode = localStorage.getItem("mode") || ""; document.documentElement.dataset.theme = localStorage.getItem("theme") || ""; </script>  <link href="../../_static/styles/theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="../../_static/styles/bootstrap.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="../../_static/styles/pydata-sphinx-theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="../../_static/vendor/fontawesome/6.5.2/css/all.min.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="../../_static/vendor/fontawesome/6.5.2/webfonts/fa-solid-900.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="../../_static/vendor/fontawesome/6.5.2/webfonts/fa-brands-400.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="../../_static/vendor/fontawesome/6.5.2/webfonts/fa-regular-400.woff2" /> <link rel="stylesheet" type="text/css" href="../../_static/pygments.css" /> <link rel="stylesheet" type="text/css" href="../../_static/styles/sphinx-book-theme.css" /> <link rel="stylesheet" type="text/css" href="../../_static/mystnb.4510f1fc1dee50b3e5859aac5469c37c29e427902b24a333a5f9fcb2f0b3ac41.css" /> <link rel="stylesheet" type="text/css" href="../../_static/sphinx-design.5ea377869091fd0449014c60fc090103.min.css" /> <link rel="stylesheet" type="text/css" href="../../_static/css/flax_theme.css" />  <link rel="preload" as="script" href="../../_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b" /> <link rel="preload" as="script" href="../../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b" /> <script src="../../_static/vendor/fontawesome/6.5.2/js/all.min.js?digest=dfe6caa3a7d634c4db9b"></script> <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script> <script src="../../_static/jquery.js"></script> <script src="../../_static/underscore.js"></script> <script src="../../_static/_sphinx_javascript_frameworks_compat.js"></script> <script src="../../_static/doctools.js"></script> <script src="../../_static/sphinx_highlight.js"></script> <script src="../../_static/scripts/sphinx-book-theme.js"></script> <script src="../../_static/design-tabs.js"></script> <script>DOCUMENTATION_OPTIONS.pagename = 'guides/converting_and_upgrading/haiku_migration_guide';</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="Convert PyTorch models to Flax" href="convert_pytorch_to_flax.html" /> <link rel="prev" title="Converting and upgrading" href="index.html" /> <meta name="viewport" content="width=device-width, initial-scale=1"/> <meta name="docsearch:language" content="en"/> <script async type="text/javascript" src="/_/static/javascript/readthedocs-addons.js"></script><meta name="readthedocs-project-slug" content="flax-linen" /><meta name="readthedocs-version-slug" content="latest" /><meta name="readthedocs-resolver-filename" content="/guides/converting_and_upgrading/haiku_migration_guide.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="../flax_fundamentals/flax_basics.html">Flax Basics</a></li> <li class="toctree-l1 current active has-children"><a class="reference internal" href="../index.html">Guides</a><details open="open"><summary></summary><ul class="current"> <li class="toctree-l2 has-children"><a class="reference internal" href="../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="../flax_fundamentals/flax_basics.html">Flax Basics</a></li> <li class="toctree-l3"><a class="reference internal" href="../flax_fundamentals/state_params.html">Managing Parameters and State</a></li> <li class="toctree-l3"><a class="reference internal" href="../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="../flax_fundamentals/arguments.html">Dealing with Flax Module arguments</a></li> <li class="toctree-l3"><a class="reference internal" href="../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="../data_preprocessing/index.html">Data preprocessing</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../data_preprocessing/full_eval.html">Processing the entire Dataset</a></li> <li class="toctree-l3"><a class="reference internal" href="../data_preprocessing/loading_datasets.html">Loading datasets</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../training_techniques/index.html">Training techniques</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../training_techniques/batch_norm.html">Batch normalization</a></li> <li class="toctree-l3"><a class="reference internal" href="../training_techniques/dropout.html">Dropout</a></li> <li class="toctree-l3"><a class="reference internal" href="../training_techniques/lr_schedule.html">Learning rate scheduling</a></li> <li class="toctree-l3"><a class="reference internal" href="../training_techniques/transfer_learning.html">Transfer learning</a></li> <li class="toctree-l3"><a class="reference internal" href="../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="../parallel_training/index.html">Parallel training</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../parallel_training/ensembling.html">Ensembling on multiple devices</a></li> <li class="toctree-l3"><a class="reference internal" href="../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="../model_inspection/index.html">Model inspection</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../model_inspection/model_surgery.html">Model surgery</a></li> <li class="toctree-l3"><a class="reference internal" href="../model_inspection/extracting_intermediates.html">Extracting intermediate values</a></li> </ul> </details></li> <li class="toctree-l2 current active has-children"><a class="reference internal" href="index.html">Converting and upgrading</a><details open="open"><summary></summary><ul class="current"> <li class="toctree-l3 current active"><a class="current reference internal" href="#">Migrating from Haiku to Flax</a></li> <li class="toctree-l3"><a class="reference internal" href="convert_pytorch_to_flax.html">Convert PyTorch models to Flax</a></li> <li class="toctree-l3"><a class="reference internal" href="orbax_upgrade_guide.html">Migrate checkpointing to Orbax</a></li> <li class="toctree-l3"><a class="reference internal" href="optax_update_guide.html">Upgrading my codebase to Optax</a></li> <li class="toctree-l3"><a class="reference internal" href="linen_upgrade_guide.html">Upgrading my codebase to Linen</a></li> <li class="toctree-l3"><a class="reference internal" href="rnncell_upgrade_guide.html">RNNCellBase Upgrade Guide</a></li> <li class="toctree-l3"><a class="reference internal" href="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="../quantization/index.html">Quantization</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../quantization/fp8_basics.html">User Guide on Using FP8</a></li> </ul> </details></li> <li class="toctree-l2"><a class="reference internal" href="../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 has-children"><a class="reference internal" href="../../api_reference/index.html">API Reference</a><details><summary></summary><ul> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.config.html">flax.config package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.core.frozen_dict.html">flax.core.frozen_dict package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.cursor.html">flax.cursor package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.errors.html">flax.errors package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.jax_utils.html">flax.jax_utils package</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../../api_reference/flax.linen/index.html">flax.linen</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/module.html">Module</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/init_apply.html">Init/Apply</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/layers.html">Layers</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/activation_functions.html">Activation functions</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/initializers.html">Initializers</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/transformations.html">Transformations</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/inspection.html">Inspection</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/variable.html">Variable dictionary</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/spmd.html">SPMD</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/decorators.html">Decorators</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/profiling.html">Profiling</a></li> </ul> </details></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.serialization.html">flax.serialization package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.struct.html">flax.struct package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.traceback_util.html">flax.traceback_util package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.training.html">flax.training package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/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/guides/converting_and_upgrading/haiku_migration_guide.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>Migrating from Haiku to Flax</h1>  <div id="print-main-content"> <div id="jb-print-toc"> <div> <h2> Contents </h2> </div> <nav aria-label="Page"> <ul class="visible nav section-nav flex-column"> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#basic-example">Basic Example</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#handling-state">Handling State</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#using-multiple-methods">Using Multiple Methods</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#lifted-transforms">Lifted Transforms</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#scan-over-layers">Scan over layers</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#top-level-haiku-functions-vs-top-level-flax-modules">Top-level Haiku functions vs top-level Flax modules</a></li> </ul> </nav> </div> </div> </div> <div id="searchbox"></div> <article class="bd-article"> <div class="section" id="migrating-from-haiku-to-flax"> <h1>Migrating from Haiku to Flax<a class="headerlink" href="#migrating-from-haiku-to-flax" title="Permalink to this heading">#</a></h1> This guide will walk through the process of migrating Haiku models to Flax, and highlight the differences between the two libraries. <div class="section" id="basic-example"> <h2>Basic Example<a class="headerlink" href="#basic-example" title="Permalink to this heading">#</a></h2> To create custom Modules you subclass from a <code class="docutils literal notranslate">Module</code> base class in both Haiku and Flax. However, Haiku classes use a regular <code class="docutils literal notranslate">__init__</code> method whereas Flax classes are <code class="docutils literal notranslate">dataclasses</code>, meaning you define some class attributes that are used to automatically generate a constructor. Also, all Flax Modules accept a <code class="docutils literal notranslate">name</code> argument without needing to define it, whereas in Haiku <code class="docutils literal notranslate">name</code> must be explicitly defined in the constructor signature and passed to the superclass constructor. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-0" name="sd-tab-set-0" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-0"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>import haiku as hk class Block(hk.Module): def __init__(self, features: int, name=None): super().__init__(name=name) self.features = features def __call__(self, x, training: bool): x = hk.Linear(self.features)(x) x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x) x = jax.nn.relu(x) return x class Model(hk.Module): def __init__(self, dmid: int, dout: int, name=None): super().__init__(name=name) self.dmid = dmid self.dout = dout def __call__(self, x, training: bool): x = Block(self.dmid)(x, training) x = hk.Linear(self.dout)(x) return x </pre></div> </div> </div> <input id="sd-tab-item-1" name="sd-tab-set-0" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-1"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>import flax.linen as nn class Block(nn.Module): features: int @nn.compact def __call__(self, x, training: bool): x = nn.Dense(self.features)(x) x = nn.Dropout(0.5, deterministic=not training)(x) x = jax.nn.relu(x) return x class Model(nn.Module): dmid: int dout: int @nn.compact def __call__(self, x, training: bool): x = Block(self.dmid)(x, training) x = nn.Dense(self.dout)(x) return x </pre></div> </div> </div> </div> The <code class="docutils literal notranslate">__call__</code> method looks very similar in both libraries, however, in Flax you have to use the <code class="docutils literal notranslate">@nn.compact</code> decorator in order to be able to define submodules inline. In Haiku, this is the default behavior. Now, a place where Haiku and Flax differ substantially is in how you construct the model. In Haiku, you use <code class="docutils literal notranslate">hk.transform</code> over a function that calls your Module, <code class="docutils literal notranslate">transform</code> will return an object with <code class="docutils literal notranslate">init</code> and <code class="docutils literal notranslate">apply</code> methods. In Flax, you simply instantiate your Module. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-2" name="sd-tab-set-1" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-2"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def forward(x, training: bool): return Model(256, 10)(x, training) model = hk.transform(forward) </pre></div> </div> </div> <input id="sd-tab-item-3" name="sd-tab-set-1" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-3"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>... model = Model(256, 10) </pre></div> </div> </div> </div> To get the model parameters in both libraries you use the <code class="docutils literal notranslate">init</code> method with a <code class="docutils literal notranslate">random.key</code> plus some inputs to run the model. The main difference here is that Flax returns a mapping from collection names to nested array dictionaries, <code class="docutils literal notranslate">params</code> is just one of these possible collections. In Haiku, you get the <code class="docutils literal notranslate">params</code> structure directly. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-4" name="sd-tab-set-2" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-4"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>sample_x = jax.numpy.ones((1, 784)) params = model.init( random.key(0), sample_x, training=False # <== inputs ) ... </pre></div> </div> </div> <input id="sd-tab-item-5" name="sd-tab-set-2" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-5"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>sample_x = jax.numpy.ones((1, 784)) variables = model.init( random.key(0), sample_x, training=False # <== inputs ) params = variables["params"] </pre></div> </div> </div> </div> One very important thing to note is that in Flax the parameters structure is hierarchical, with one level per nested module and a final level for the parameter name. In Haiku the parameters structure is a python dictionary with a two level hierarchy: the fully qualified module name mapping to the parameter name. The module name consists of a <code class="docutils literal notranslate">/</code> separated string path of all the nested Modules. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-6" name="sd-tab-set-3" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-6"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>... { 'model/block/linear': { 'b': (256,), 'w': (784, 256), }, 'model/linear': { 'b': (10,), 'w': (256, 10), } } ... </pre></div> </div> </div> <input id="sd-tab-item-7" name="sd-tab-set-3" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-7"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>FrozenDict({ Block_0: { Dense_0: { bias: (256,), kernel: (784, 256), }, }, Dense_0: { bias: (10,), kernel: (256, 10), }, }) </pre></div> </div> </div> </div> During training in both frameworks you pass the parameters structure to the <code class="docutils literal notranslate">apply</code> method to run the forward pass. Since we are using dropout, in both cases we must provide a <code class="docutils literal notranslate">key</code> to <code class="docutils literal notranslate">apply</code> in order to generate the random dropout masks. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-8" name="sd-tab-set-4" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-8"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def train_step(key, params, inputs, labels): def loss_fn(params): logits = model.apply( params, key, inputs, training=True # <== inputs ) return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean() grads = jax.grad(loss_fn)(params) params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads) return params </pre></div> </div> </div> <input id="sd-tab-item-9" name="sd-tab-set-4" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-9"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def train_step(key, params, inputs, labels): def loss_fn(params): logits = model.apply( {'params': params}, inputs, training=True, # <== inputs rngs={'dropout': key} ) return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean() grads = jax.grad(loss_fn)(params) params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads) return params </pre></div> </div> </div> </div> The most notable differences is that in Flax you have to pass the parameters inside a dictionary with a <code class="docutils literal notranslate">params</code> key, and the key inside a dictionary with a <code class="docutils literal notranslate">dropout</code> key. This is because in Flax you can have many types of model state and random state. In Haiku, you just pass the parameters and the key directly. </div> <div class="section" id="handling-state"> <h2>Handling State<a class="headerlink" href="#handling-state" title="Permalink to this heading">#</a></h2> Now let’s see how mutable state is handled in both libraries. We will take the same model as before, but now we will replace Dropout with BatchNorm. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-10" name="sd-tab-set-5" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-10"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class Block(hk.Module): def __init__(self, features: int, name=None): super().__init__(name=name) self.features = features def __call__(self, x, training: bool): x = hk.Linear(self.features)(x) x = hk.BatchNorm( create_scale=True, create_offset=True, decay_rate=0.99 )(x, is_training=training) x = jax.nn.relu(x) return x </pre></div> </div> </div> <input id="sd-tab-item-11" name="sd-tab-set-5" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-11"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class Block(nn.Module): features: int @nn.compact def __call__(self, x, training: bool): x = nn.Dense(self.features)(x) x = nn.BatchNorm( momentum=0.99 )(x, use_running_average=not training) x = jax.nn.relu(x) return x </pre></div> </div> </div> </div> The code is very similar in this case as both libraries provide a BatchNorm layer. The most notable difference is that Haiku uses <code class="docutils literal notranslate">is_training</code> to control whether or not to update the running statistics, whereas Flax uses <code class="docutils literal notranslate">use_running_average</code> for the same purpose. To instantiate a stateful model in Haiku you use <code class="docutils literal notranslate">hk.transform_with_state</code>, which changes the signature for <code class="docutils literal notranslate">init</code> and <code class="docutils literal notranslate">apply</code> to accept and return state. As before, in Flax you construct the Module directly. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-12" name="sd-tab-set-6" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-12"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def forward(x, training: bool): return Model(256, 10)(x, training) model = hk.transform_with_state(forward) </pre></div> </div> </div> <input id="sd-tab-item-13" name="sd-tab-set-6" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-13"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>... model = Model(256, 10) </pre></div> </div> </div> </div> To initialize both the parameters and state you just call the <code class="docutils literal notranslate">init</code> method as before. However, in Haiku you now get <code class="docutils literal notranslate">state</code> as a second return value, and in Flax you get a new <code class="docutils literal notranslate">batch_stats</code> collection in the <code class="docutils literal notranslate">variables</code> dictionary. Note that since <code class="docutils literal notranslate">hk.BatchNorm</code> only initializes batch statistics when <code class="docutils literal notranslate">is_training=True</code>, we must set <code class="docutils literal notranslate">training=True</code> when initializing parameters of a Haiku model with an <code class="docutils literal notranslate">hk.BatchNorm</code> layer. In Flax, we can set <code class="docutils literal notranslate">training=False</code> as usual. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-14" name="sd-tab-set-7" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-14"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>sample_x = jax.numpy.ones((1, 784)) params, state = model.init( random.key(0), sample_x, training=True # <== inputs ) ... </pre></div> </div> </div> <input id="sd-tab-item-15" name="sd-tab-set-7" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-15"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>sample_x = jax.numpy.ones((1, 784)) variables = model.init( random.key(0), sample_x, training=False # <== inputs ) params, batch_stats = variables["params"], variables["batch_stats"] </pre></div> </div> </div> </div> In general, in Flax you might find other state collections in the <code class="docutils literal notranslate">variables</code> dictionary such as <code class="docutils literal notranslate">cache</code> for auto-regressive transformers models, <code class="docutils literal notranslate">intermediates</code> for intermediate values added using <code class="docutils literal notranslate">Module.sow</code>, or other collection names defined by custom layers. Haiku only makes a distinction between <code class="docutils literal notranslate">params</code> (variables which do not change while running <code class="docutils literal notranslate">apply</code>) and <code class="docutils literal notranslate">state</code> (variables which can change while running <code class="docutils literal notranslate">apply</code>). Now, training looks very similar in both frameworks as you use the same <code class="docutils literal notranslate">apply</code> method to run the forward pass. In Haiku, now pass the <code class="docutils literal notranslate">state</code> as the second argument to <code class="docutils literal notranslate">apply</code>, and get the new state as the second return value. In Flax, you instead add <code class="docutils literal notranslate">batch_stats</code> as a new key to the input dictionary, and get the <code class="docutils literal notranslate">updates</code> variables dictionary as the second return value. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-16" name="sd-tab-set-8" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-16"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def train_step(params, state, inputs, labels): def loss_fn(params): logits, new_state = model.apply( params, state, None, # <== rng inputs, training=True # <== inputs ) loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean() return loss, new_state grads, new_state = jax.grad(loss_fn, has_aux=True)(params) params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads) return params, new_state </pre></div> </div> </div> <input id="sd-tab-item-17" name="sd-tab-set-8" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-17"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def train_step(params, batch_stats, inputs, labels): def loss_fn(params): logits, updates = model.apply( {'params': params, 'batch_stats': batch_stats}, inputs, training=True, # <== inputs mutable='batch_stats', ) loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean() return loss, updates["batch_stats"] grads, batch_stats = jax.grad(loss_fn, has_aux=True)(params) params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads) return params, batch_stats </pre></div> </div> </div> </div> One major difference is that in Flax a state collection can be mutable or immutable. During <code class="docutils literal notranslate">init</code> all collections are mutable by default, however, during <code class="docutils literal notranslate">apply</code> you have to explicitly specify which collections are mutable. In this example, we specify that <code class="docutils literal notranslate">batch_stats</code> is mutable. Here a single string is passed but a list can also be given if there are more mutable collections. If this is not done an error will be raised at runtime when trying to mutate <code class="docutils literal notranslate">batch_stats</code>. Also, when <code class="docutils literal notranslate">mutable</code> is anything other than <code class="docutils literal notranslate">False</code>, the <code class="docutils literal notranslate">updates</code> dictionary is returned as the second return value of <code class="docutils literal notranslate">apply</code>, else only the model output is returned. Haiku makes the mutable/immutable distinction through having <code class="docutils literal notranslate">params</code> (immutable) and <code class="docutils literal notranslate">state</code> (mutable) and using either <code class="docutils literal notranslate">hk.transform</code> or <code class="docutils literal notranslate">hk.transform_with_state</code> </div> <div class="section" id="using-multiple-methods"> <h2>Using Multiple Methods<a class="headerlink" href="#using-multiple-methods" title="Permalink to this heading">#</a></h2> In this section we will take a look at how to use multiple methods in Haiku and Flax. As an example, we will implement an auto-encoder model with three methods: <code class="docutils literal notranslate">encode</code>, <code class="docutils literal notranslate">decode</code>, and <code class="docutils literal notranslate">__call__</code>. In Haiku, we can just define the submodules that <code class="docutils literal notranslate">encode</code> and <code class="docutils literal notranslate">decode</code> need directly in <code class="docutils literal notranslate">__init__</code>, in this case each will just use a <code class="docutils literal notranslate">Linear</code> layer. In Flax, we will define an <code class="docutils literal notranslate">encoder</code> and a <code class="docutils literal notranslate">decoder</code> Module ahead of time in <code class="docutils literal notranslate">setup</code>, and use them in the <code class="docutils literal notranslate">encode</code> and <code class="docutils literal notranslate">decode</code> respectively. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-18" name="sd-tab-set-9" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-18"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class AutoEncoder(hk.Module): def __init__(self, embed_dim: int, output_dim: int, name=None): super().__init__(name=name) self.encoder = hk.Linear(embed_dim, name="encoder") self.decoder = hk.Linear(output_dim, name="decoder") def encode(self, x): return self.encoder(x) def decode(self, x): return self.decoder(x) def __call__(self, x): x = self.encode(x) x = self.decode(x) return x </pre></div> </div> </div> <input id="sd-tab-item-19" name="sd-tab-set-9" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-19"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class AutoEncoder(nn.Module): embed_dim: int output_dim: int def setup(self): self.encoder = nn.Dense(self.embed_dim) self.decoder = nn.Dense(self.output_dim) def encode(self, x): return self.encoder(x) def decode(self, x): return self.decoder(x) def __call__(self, x): x = self.encode(x) x = self.decode(x) return x </pre></div> </div> </div> </div> Note that in Flax <code class="docutils literal notranslate">setup</code> doesn’t run after <code class="docutils literal notranslate">__init__</code>, instead it runs when <code class="docutils literal notranslate">init</code> or <code class="docutils literal notranslate">apply</code> are called. Now, we want to be able to call any method from our <code class="docutils literal notranslate">AutoEncoder</code> model. In Haiku we can define multiple <code class="docutils literal notranslate">apply</code> methods for a module through <code class="docutils literal notranslate">hk.multi_transform</code>. The function passed to <code class="docutils literal notranslate">multi_transform</code> defines how to initialize the module and which different apply methods to generate. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-20" name="sd-tab-set-10" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-20"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def forward(): module = AutoEncoder(256, 784) init = lambda x: module(x) return init, (module.encode, module.decode) model = hk.multi_transform(forward) </pre></div> </div> </div> <input id="sd-tab-item-21" name="sd-tab-set-10" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-21"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>... model = AutoEncoder(256, 784) </pre></div> </div> </div> </div> To initialize the parameters of our model, <code class="docutils literal notranslate">init</code> can be used to trigger the <code class="docutils literal notranslate">__call__</code> method, which uses both the <code class="docutils literal notranslate">encode</code> and <code class="docutils literal notranslate">decode</code> method. This will create all the necessary parameters for the model. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-22" name="sd-tab-set-11" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-22"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>params = model.init( random.key(0), x=jax.numpy.ones((1, 784)), ) ... </pre></div> </div> </div> <input id="sd-tab-item-23" name="sd-tab-set-11" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-23"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>variables = model.init( random.key(0), x=jax.numpy.ones((1, 784)), ) params = variables["params"] </pre></div> </div> </div> </div> This generates the following parameter structure. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-24" name="sd-tab-set-12" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-24"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>{ 'auto_encoder/~/decoder': { 'b': (784,), 'w': (256, 784) }, 'auto_encoder/~/encoder': { 'b': (256,), 'w': (784, 256) } } </pre></div> </div> </div> <input id="sd-tab-item-25" name="sd-tab-set-12" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-25"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>FrozenDict({ decoder: { bias: (784,), kernel: (256, 784), }, encoder: { bias: (256,), kernel: (784, 256), }, }) </pre></div> </div> </div> </div> Finally, let’s explore how we can employ the <code class="docutils literal notranslate">apply</code> function to invoke the <code class="docutils literal notranslate">encode</code> method: <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-26" name="sd-tab-set-13" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-26"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>encode, decode = model.apply z = encode( params, None, # <== rng x=jax.numpy.ones((1, 784)), ) </pre></div> </div> </div> <input id="sd-tab-item-27" name="sd-tab-set-13" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-27"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>... z = model.apply( {"params": params}, x=jax.numpy.ones((1, 784)), method="encode", ) </pre></div> </div> </div> </div> Because the Haiku <code class="docutils literal notranslate">apply</code> function is generated through <code class="docutils literal notranslate">hk.multi_transform</code>, it’s a tuple of two functions which we can unpack into an <code class="docutils literal notranslate">encode</code> and <code class="docutils literal notranslate">decode</code> function which correspond to the methods on the <code class="docutils literal notranslate">AutoEncoder</code> module. In Flax we call the <code class="docutils literal notranslate">encode</code> method through passing the method name as a string. Another noteworthy distinction here is that in Haiku, <code class="docutils literal notranslate">rng</code> needs to be explicitly passed, even though the module does not use any stochastic operations during <code class="docutils literal notranslate">apply</code>. In Flax this is not necessary (check out <a class="reference external" href="https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/rng_guide.html">Randomness and PRNGs in Flax</a>). The Haiku <code class="docutils literal notranslate">rng</code> is set to <code class="docutils literal notranslate">None</code> here, but you could also use <code class="docutils literal notranslate">hk.without_apply_rng</code> on the <code class="docutils literal notranslate">apply</code> function to remove the <code class="docutils literal notranslate">rng</code> argument. </div> <div class="section" id="lifted-transforms"> <h2>Lifted Transforms<a class="headerlink" href="#lifted-transforms" title="Permalink to this heading">#</a></h2> Both Flax and Haiku provide a set of transforms, which we will refer to as lifted transforms, that wrap JAX transformations in such a way that they can be used with Modules and sometimes provide additional functionality. In this section we will take a look at how to use the lifted version of <code class="docutils literal notranslate">scan</code> in both Flax and Haiku to implement a simple RNN layer. To begin, we will first define a <code class="docutils literal notranslate">RNNCell</code> module that will contain the logic for a single step of the RNN. We will also define a <code class="docutils literal notranslate">initial_state</code> method that will be used to initialize the state (a.k.a. <code class="docutils literal notranslate">carry</code>) of the RNN. Like with <code class="docutils literal notranslate">jax.lax.scan</code>, the <code class="docutils literal notranslate">RNNCell.__call__</code> method will be a function that takes the carry and input, and returns the new carry and output. In this case, the carry and the output are the same. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-28" name="sd-tab-set-14" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-28"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class RNNCell(hk.Module): def __init__(self, hidden_size: int, name=None): super().__init__(name=name) self.hidden_size = hidden_size def __call__(self, carry, x): x = jnp.concatenate([carry, x], axis=-1) x = hk.Linear(self.hidden_size)(x) x = jax.nn.relu(x) return x, x def initial_state(self, batch_size: int): return jnp.zeros((batch_size, self.hidden_size)) </pre></div> </div> </div> <input id="sd-tab-item-29" name="sd-tab-set-14" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-29"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class RNNCell(nn.Module): hidden_size: int @nn.compact def __call__(self, carry, x): x = jnp.concatenate([carry, x], axis=-1) x = nn.Dense(self.hidden_size)(x) x = jax.nn.relu(x) return x, x def initial_state(self, batch_size: int): return jnp.zeros((batch_size, self.hidden_size)) </pre></div> </div> </div> </div> Next, we will define a <code class="docutils literal notranslate">RNN</code> Module that will contain the logic for the entire RNN. In Haiku, we will first initialze the <code class="docutils literal notranslate">RNNCell</code>, then use it to construct the <code class="docutils literal notranslate">carry</code>, and finally use <code class="docutils literal notranslate">hk.scan</code> to run the <code class="docutils literal notranslate">RNNCell</code> over the input sequence. In Flax its done a bit differently, we will use <code class="docutils literal notranslate">nn.scan</code> to define a new temporary type that wraps <code class="docutils literal notranslate">RNNCell</code>. During this process we will also specify instruct <code class="docutils literal notranslate">nn.scan</code> to broadcast the <code class="docutils literal notranslate">params</code> collection (all steps share the same parameters) and to not split the <code class="docutils literal notranslate">params</code> rng stream (so all steps intialize with the same parameters), and finally we will specify that we want scan to run over the second axis of the input and stack the outputs along the second axis as well. We will then use this temporary type immediately to create an instance of the lifted <code class="docutils literal notranslate">RNNCell</code> and use it to create the <code class="docutils literal notranslate">carry</code> and the run the <code class="docutils literal notranslate">__call__</code> method which will <code class="docutils literal notranslate">scan</code> over the sequence. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-30" name="sd-tab-set-15" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-30"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class RNN(hk.Module): def __init__(self, hidden_size: int, name=None): super().__init__(name=name) self.hidden_size = hidden_size def __call__(self, x): cell = RNNCell(self.hidden_size) carry = cell.initial_state(x.shape[0]) carry, y = hk.scan(cell, carry, jnp.swapaxes(x, 1, 0)) y = jnp.swapaxes(y, 0, 1) return y </pre></div> </div> </div> <input id="sd-tab-item-31" name="sd-tab-set-15" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-31"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class RNN(nn.Module): hidden_size: int @nn.compact def __call__(self, x): rnn = nn.scan(RNNCell, variable_broadcast='params', split_rngs={'params': False}, in_axes=1, out_axes=1)(self.hidden_size) carry = rnn.initial_state(x.shape[0]) carry, y = rnn(carry, x) return y </pre></div> </div> </div> </div> In general, the main difference between lifted transforms between Flax and Haiku is that in Haiku the lifted transforms don’t operate over the state, that is, Haiku will handle the <code class="docutils literal notranslate">params</code> and <code class="docutils literal notranslate">state</code> in such a way that it keeps the same shape inside and outside of the transform. In Flax, the lifted transforms can operate over both variable collections and rng streams, the user must define how different collections are treated by each transform according to the transform’s semantics. Finally, let’s quickly view how the <code class="docutils literal notranslate">RNN</code> Module would be used in both Haiku and Flax. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-32" name="sd-tab-set-16" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-32"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def forward(x): return RNN(64)(x) model = hk.without_apply_rng(hk.transform(forward)) params = model.init( random.key(0), x=jax.numpy.ones((3, 12, 32)), ) y = model.apply( params, x=jax.numpy.ones((3, 12, 32)), ) </pre></div> </div> </div> <input id="sd-tab-item-33" name="sd-tab-set-16" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-33"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>... model = RNN(64) variables = model.init( random.key(0), x=jax.numpy.ones((3, 12, 32)), ) params = variables['params'] y = model.apply( {'params': params}, x=jax.numpy.ones((3, 12, 32)), ) </pre></div> </div> </div> </div> The only notable change with respect to the examples in the previous sections is that this time around we used <code class="docutils literal notranslate">hk.without_apply_rng</code> in Haiku so we didn’t have to pass the <code class="docutils literal notranslate">rng</code> argument as <code class="docutils literal notranslate">None</code> to the <code class="docutils literal notranslate">apply</code> method. </div> <div class="section" id="scan-over-layers"> <h2>Scan over layers<a class="headerlink" href="#scan-over-layers" title="Permalink to this heading">#</a></h2> One very important application of <code class="docutils literal notranslate">scan</code> is apply a sequence of layers iteratively over an input, passing the output of each layer as the input to the next layer. This is very useful to reduce compilation time for big models. As an example we will create a simple <code class="docutils literal notranslate">Block</code> Module, and then use it inside an <code class="docutils literal notranslate">MLP</code> Module that will apply the <code class="docutils literal notranslate">Block</code> Module <code class="docutils literal notranslate">num_layers</code> times. In Haiku, we define the <code class="docutils literal notranslate">Block</code> Module as usual, and then inside <code class="docutils literal notranslate">MLP</code> we will use <code class="docutils literal notranslate">hk.experimental.layer_stack</code> over a <code class="docutils literal notranslate">stack_block</code> function to create a stack of <code class="docutils literal notranslate">Block</code> Modules. In Flax, the definition of <code class="docutils literal notranslate">Block</code> is a little different, <code class="docutils literal notranslate">__call__</code> will accept and return a second dummy input/output that in both cases will be <code class="docutils literal notranslate">None</code>. In <code class="docutils literal notranslate">MLP</code>, we will use <code class="docutils literal notranslate">nn.scan</code> as in the previous example, but by setting <code class="docutils literal notranslate">split_rngs={'params': True}</code> and <code class="docutils literal notranslate">variable_axes={'params': 0}</code> we are telling <code class="docutils literal notranslate">nn.scan</code> create different parameters for each step and slice the <code class="docutils literal notranslate">params</code> collection along the first axis, effectively implementing a stack of <code class="docutils literal notranslate">Block</code> Modules as in Haiku. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-34" name="sd-tab-set-17" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-34"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class Block(hk.Module): def __init__(self, features: int, name=None): super().__init__(name=name) self.features = features def __call__(self, x, training: bool): x = hk.Linear(self.features)(x) x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x) x = jax.nn.relu(x) return x class MLP(hk.Module): def __init__(self, features: int, num_layers: int, name=None): super().__init__(name=name) self.features = features self.num_layers = num_layers def __call__(self, x, training: bool): @hk.experimental.layer_stack(self.num_layers) def stack_block(x): return Block(self.features)(x, training) stack = hk.experimental.layer_stack(self.num_layers) return stack_block(x) </pre></div> </div> </div> <input id="sd-tab-item-35" name="sd-tab-set-17" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-35"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class Block(nn.Module): features: int training: bool @nn.compact def __call__(self, x, _): x = nn.Dense(self.features)(x) x = nn.Dropout(0.5)(x, deterministic=not self.training) x = jax.nn.relu(x) return x, None class MLP(nn.Module): features: int num_layers: int @nn.compact def __call__(self, x, training: bool): ScanBlock = nn.scan( Block, variable_axes={'params': 0}, split_rngs={'params': True}, length=self.num_layers) y, _ = ScanBlock(self.features, training)(x, None) return y </pre></div> </div> </div> </div> Notice how in Flax we pass <code class="docutils literal notranslate">None</code> as the second argument to <code class="docutils literal notranslate">ScanBlock</code> and ignore its second output. These represent the inputs/outputs per-step but they are <code class="docutils literal notranslate">None</code> because in this case we don’t have any. Initializing each model is the same as in previous examples. In this case, we will be specifying that we want to use <code class="docutils literal notranslate">5</code> layers each with <code class="docutils literal notranslate">64</code> features. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-36" name="sd-tab-set-18" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-36"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def forward(x, training: bool): return MLP(64, num_layers=5)(x, training) model = hk.transform(forward) sample_x = jax.numpy.ones((1, 64)) params = model.init( random.key(0), sample_x, training=False # <== inputs ) ... </pre></div> </div> </div> <input id="sd-tab-item-37" name="sd-tab-set-18" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-37"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>... model = MLP(64, num_layers=5) sample_x = jax.numpy.ones((1, 64)) variables = model.init( random.key(0), sample_x, training=False # <== inputs ) params = variables['params'] </pre></div> </div> </div> </div> When using scan over layers the one thing you should notice is that all layers are fused into a single layer whose parameters have an extra “layer” dimension on the first axis. In this case, the shape of all parameters will start with <code class="docutils literal notranslate">(5, ...)</code> as we are using <code class="docutils literal notranslate">5</code> layers. <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-38" name="sd-tab-set-19" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-38"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>... { 'mlp/__layer_stack_no_per_layer/block/linear': { 'b': (5, 64), 'w': (5, 64, 64) } } ... </pre></div> </div> </div> <input id="sd-tab-item-39" name="sd-tab-set-19" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-39"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>FrozenDict({ ScanBlock_0: { Dense_0: { bias: (5, 64), kernel: (5, 64, 64), }, }, }) </pre></div> </div> </div> </div> </div> <div class="section" id="top-level-haiku-functions-vs-top-level-flax-modules"> <h2>Top-level Haiku functions vs top-level Flax modules<a class="headerlink" href="#top-level-haiku-functions-vs-top-level-flax-modules" title="Permalink to this heading">#</a></h2> In Haiku, it is possible to write the entire model as a single function by using the raw <code class="docutils literal notranslate">hk.{get,set}_{parameter,state}</code> to define/access model parameters and states. It very common to write the top-level “Module” as a function instead: The Flax team recommends a more Module-centric approach that uses <cite>__call__</cite> to define the forward function. The corresponding accessor will be <cite>nn.module.param</cite> and <cite>nn.module.variable</cite> (go to <a class="reference external" href="#handling-state">Handling State</a> for an explanaion on collections). <div class="sd-tab-set docutils"> <input checked="checked" id="sd-tab-item-40" name="sd-tab-set-20" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Haiku" for="sd-tab-item-40"> Haiku</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>def forward(x): counter = hk.get_state('counter', shape=[], dtype=jnp.int32, init=jnp.ones) multiplier = hk.get_parameter('multiplier', shape=[1,], dtype=x.dtype, init=jnp.ones) output = x + multiplier * counter hk.set_state("counter", counter + 1) return output model = hk.transform_with_state(forward) params, state = model.init(random.key(0), jax.numpy.ones((1, 64))) </pre></div> </div> </div> <input id="sd-tab-item-41" name="sd-tab-set-20" type="radio"> <label class="sd-tab-label" data-sync-group="tab" data-sync-id="Flax" for="sd-tab-item-41"> Flax</label><div class="sd-tab-content docutils"> <div class="highlight-python notranslate"><div class="highlight"><pre>class FooModule(nn.Module): @nn.compact def __call__(self, x): counter = self.variable('counter', 'count', lambda: jnp.ones((), jnp.int32)) multiplier = self.param('multiplier', nn.initializers.ones_init(), [1,], x.dtype) output = x + multiplier * counter.value if not self.is_initializing(): # otherwise model.init() also increases it counter.value += 1 return output model = FooModule() variables = model.init(random.key(0), jax.numpy.ones((1, 64))) params, counter = variables['params'], variables['counter'] </pre></div> </div> </div> </div> </div> </div> </article> <footer class="prev-next-footer d-print-none"> <div class="prev-next-area"> <a class="left-prev" href="index.html" title="previous page"> <div class="prev-next-info"> previous Converting and upgrading </div> </a> <a class="right-next" href="convert_pytorch_to_flax.html" title="next page"> <div class="prev-next-info"> next Convert PyTorch models to Flax </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="#basic-example">Basic Example</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#handling-state">Handling State</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#using-multiple-methods">Using Multiple Methods</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#lifted-transforms">Lifted Transforms</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#scan-over-layers">Scan over layers</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#top-level-haiku-functions-vs-top-level-flax-modules">Top-level Haiku functions vs top-level Flax modules</a></li> </ul> </nav></div> </div></div> </div> <footer class="bd-footer-content"> <div class="bd-footer-content__inner container"> <div class="footer-item"> By The Flax authors </div> <div class="footer-item"> © Copyright 2023, The Flax authors. </div> <div class="footer-item"> </div> <div class="footer-item"> </div> </div> </footer> </main> </div> </div>  <script src="../../_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b"></script> <script src="../../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b"></script> <footer class="bd-footer"> </footer> </body> </html>