CINXE.COM

<!DOCTYPE html> <html lang="en" data-content_root="" > <head> <meta charset="utf-8" /> <meta name="viewport" content="width=device-width, initial-scale=1.0" /> <title>flax.serialization package</title> <script data-cfasync="false"> document.documentElement.dataset.mode = localStorage.getItem("mode") || ""; document.documentElement.dataset.theme = localStorage.getItem("theme") || ""; </script>  <link href="../_static/styles/theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="../_static/styles/bootstrap.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="../_static/styles/pydata-sphinx-theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="../_static/vendor/fontawesome/6.5.2/css/all.min.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="../_static/vendor/fontawesome/6.5.2/webfonts/fa-solid-900.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="../_static/vendor/fontawesome/6.5.2/webfonts/fa-brands-400.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="../_static/vendor/fontawesome/6.5.2/webfonts/fa-regular-400.woff2" /> <link rel="stylesheet" type="text/css" href="../_static/pygments.css" /> <link rel="stylesheet" type="text/css" href="../_static/styles/sphinx-book-theme.css" /> <link rel="stylesheet" type="text/css" href="../_static/mystnb.4510f1fc1dee50b3e5859aac5469c37c29e427902b24a333a5f9fcb2f0b3ac41.css" /> <link rel="stylesheet" type="text/css" href="../_static/sphinx-design.5ea377869091fd0449014c60fc090103.min.css" /> <link rel="stylesheet" type="text/css" href="../_static/css/flax_theme.css" />  <link rel="preload" as="script" href="../_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b" /> <link rel="preload" as="script" href="../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b" /> <script src="../_static/vendor/fontawesome/6.5.2/js/all.min.js?digest=dfe6caa3a7d634c4db9b"></script> <script data-url_root="../" id="documentation_options" src="../_static/documentation_options.js"></script> <script src="../_static/jquery.js"></script> <script src="../_static/underscore.js"></script> <script src="../_static/_sphinx_javascript_frameworks_compat.js"></script> <script src="../_static/doctools.js"></script> <script src="../_static/sphinx_highlight.js"></script> <script src="../_static/scripts/sphinx-book-theme.js"></script> <script src="../_static/design-tabs.js"></script> <script>DOCUMENTATION_OPTIONS.pagename = 'api_reference/flax.serialization';</script> <link rel="shortcut icon" href="../_static/flax.png"/> <link rel="index" title="Index" href="../genindex.html" /> <link rel="search" title="Search" href="../search.html" /> <link rel="next" title="flax.struct package" href="flax.struct.html" /> <link rel="prev" title="Profiling" href="flax.linen/profiling.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.serialization.html" /><meta name="readthedocs-http-status" content="200" /></head> <body data-bs-spy="scroll" data-bs-target=".bd-toc-nav" data-offset="180" data-bs-root-margin="0px 0px -60%" data-default-mode=""> <div id="pst-skip-link" class="skip-link d-print-none"><a href="#main-content">Skip to main content</a></div> <div id="pst-scroll-pixel-helper"></div> <button type="button" class="btn rounded-pill" id="pst-back-to-top"> <i class="fa-solid fa-arrow-up"></i>Back to top</button> <input type="checkbox" class="sidebar-toggle" id="pst-primary-sidebar-checkbox"/> <label class="overlay overlay-primary" for="pst-primary-sidebar-checkbox"></label> <input type="checkbox" class="sidebar-toggle" id="pst-secondary-sidebar-checkbox"/> <label class="overlay overlay-secondary" for="pst-secondary-sidebar-checkbox"></label> <div class="search-button__wrapper"> <div class="search-button__overlay"></div> <div class="search-button__search-container"> <form class="bd-search d-flex align-items-center" action="../search.html" method="get"> <i class="fa-solid fa-magnifying-glass"></i> <input type="search" class="form-control" name="q" id="search-input" placeholder="Search..." aria-label="Search..." autocomplete="off" autocorrect="off" autocapitalize="off" spellcheck="false"/> <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd>K</kbd></span> </form></div> </div> <div class="pst-async-banner-revealer d-none"> <aside id="bd-header-version-warning" class="d-none d-print-none" aria-label="Version warning"></aside> </div> <aside class="bd-header-announcement" aria-label="Announcement"> <div class="bd-header-announcement__content"> <a href="https://flax.readthedocs.io/en/latest/index.html" style="text-decoration: none; color: white;" > This site covers the old Flax Linen API. <span style="color: lightgray;">[Explore the new <b>Flax NNX</b> API ✨]</span> </a> </div> </aside> <header class="bd-header navbar navbar-expand-lg bd-navbar d-print-none"> </header> <div class="bd-container"> <div class="bd-container__inner bd-page-width"> <div class="bd-sidebar-primary bd-sidebar"> <div class="sidebar-header-items sidebar-primary__section"> </div> <div class="sidebar-primary-items__start sidebar-primary__section"> <div class="sidebar-primary-item"> <a class="navbar-brand logo" href="../index.html"> <img src="../_static/flax.png" class="logo__image only-light" alt=" - Home"/> <script>document.write(`<img src="../_static/flax.png" class="logo__image only-dark" alt=" - Home"/>`);</script> </a></div> <div class="sidebar-primary-item"> <script> document.write(` <button class="btn search-button-field search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip"> <i class="fa-solid fa-magnifying-glass"></i> <span class="search-button__default-text">Search</span> <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd></span> </button> `); </script></div> <div class="sidebar-primary-item"><nav class="bd-links bd-docs-nav" aria-label="Main"> <div class="bd-toc-item navbar-nav active"> <ul class="current nav bd-sidenav"> <li class="toctree-l1"><a class="reference internal" href="../quick_start.html">Quick start</a></li> <li class="toctree-l1"><a class="reference internal" href="../guides/flax_fundamentals/flax_basics.html">Flax Basics</a></li> <li class="toctree-l1 has-children"><a class="reference internal" href="../guides/index.html">Guides</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l2 has-children"><a class="reference internal" href="../guides/flax_fundamentals/index.html">Flax fundamentals</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l3"><a class="reference external" href="https://jax.readthedocs.io/en/latest/jax-101/index.html">JAX 101</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/flax_fundamentals/flax_basics.html">Flax Basics</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/flax_fundamentals/state_params.html">Managing Parameters and State</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/flax_fundamentals/setup_or_nncompact.html"><code class="docutils literal notranslate"><span class="pre">setup</span></code> vs <code class="docutils literal notranslate"><span class="pre">compact</span></code></a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/flax_fundamentals/arguments.html">Dealing with Flax Module arguments</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/flax_fundamentals/rng_guide.html">Randomness and PRNGs in Flax</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../guides/data_preprocessing/index.html">Data preprocessing</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../guides/data_preprocessing/full_eval.html">Processing the entire Dataset</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/data_preprocessing/loading_datasets.html">Loading datasets</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../guides/training_techniques/index.html">Training techniques</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../guides/training_techniques/batch_norm.html">Batch normalization</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/training_techniques/dropout.html">Dropout</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/training_techniques/lr_schedule.html">Learning rate scheduling</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/training_techniques/transfer_learning.html">Transfer learning</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/training_techniques/use_checkpointing.html">Save and load checkpoints</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../guides/parallel_training/index.html">Parallel training</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../guides/parallel_training/ensembling.html">Ensembling on multiple devices</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/parallel_training/flax_on_pjit.html">Scale up Flax Modules on multiple devices</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../guides/model_inspection/index.html">Model inspection</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../guides/model_inspection/model_surgery.html">Model surgery</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/model_inspection/extracting_intermediates.html">Extracting intermediate values</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../guides/converting_and_upgrading/index.html">Converting and upgrading</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../guides/converting_and_upgrading/haiku_migration_guide.html">Migrating from Haiku to Flax</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/converting_and_upgrading/convert_pytorch_to_flax.html">Convert PyTorch models to Flax</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/converting_and_upgrading/orbax_upgrade_guide.html">Migrate checkpointing to Orbax</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/converting_and_upgrading/optax_update_guide.html">Upgrading my codebase to Optax</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/converting_and_upgrading/linen_upgrade_guide.html">Upgrading my codebase to Linen</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/converting_and_upgrading/rnncell_upgrade_guide.html">RNNCellBase Upgrade Guide</a></li> <li class="toctree-l3"><a class="reference internal" href="../guides/converting_and_upgrading/regular_dict_upgrade_guide.html">Migrate to regular dicts</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../guides/quantization/index.html">Quantization</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../guides/quantization/fp8_basics.html">User Guide on Using FP8</a></li> </ul> </details></li> <li class="toctree-l2"><a class="reference internal" href="../guides/flax_sharp_bits.html">The Sharp Bits</a></li> </ul> </details></li> <li class="toctree-l1 has-children"><a class="reference internal" href="../examples/index.html">Examples</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l2"><a class="reference internal" href="../examples/core_examples.html">Core examples</a></li> <li class="toctree-l2"><a class="reference internal" href="../examples/google_research_examples.html">Google Research examples</a></li> <li class="toctree-l2"><a class="reference internal" href="../examples/repositories_that_use_flax.html">Repositories that use Flax</a></li> <li class="toctree-l2"><a class="reference internal" href="../examples/community_examples.html">Community examples</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference internal" href="../glossary.html">Glossary</a></li> <li class="toctree-l1"><a class="reference internal" href="../faq.html">Frequently Asked Questions (FAQ)</a></li> <li class="toctree-l1 has-children"><a class="reference internal" href="../developer_notes/index.html">Developer notes</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l2"><a class="reference internal" href="../developer_notes/module_lifecycle.html">The Flax Module lifecycle</a></li> <li class="toctree-l2"><a class="reference internal" href="../developer_notes/lift.html">Lifted transformations</a></li> <li class="toctree-l2"><a class="reference external" href="https://github.com/google/flax/tree/main/docs/flip">FLIPs</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference internal" href="../philosophy.html">The Flax philosophy</a></li> <li class="toctree-l1"><a class="reference internal" href="../contributing.html">How to contribute</a></li> <li class="toctree-l1 current active has-children"><a class="reference internal" href="index.html">API Reference</a><details open="open"><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul class="current"> <li class="toctree-l2"><a class="reference internal" href="flax.config.html">flax.config package</a></li> <li class="toctree-l2"><a class="reference internal" href="flax.core.frozen_dict.html">flax.core.frozen_dict package</a></li> <li class="toctree-l2"><a class="reference internal" href="flax.cursor.html">flax.cursor package</a></li> <li class="toctree-l2"><a class="reference internal" href="flax.errors.html">flax.errors package</a></li> <li class="toctree-l2"><a class="reference internal" href="flax.jax_utils.html">flax.jax_utils package</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="flax.linen/index.html">flax.linen</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul> <li class="toctree-l3"><a class="reference internal" href="flax.linen/module.html">Module</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/init_apply.html">Init/Apply</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/layers.html">Layers</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/activation_functions.html">Activation functions</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/initializers.html">Initializers</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/transformations.html">Transformations</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/inspection.html">Inspection</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/variable.html">Variable dictionary</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/spmd.html">SPMD</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/decorators.html">Decorators</a></li> <li class="toctree-l3"><a class="reference internal" href="flax.linen/profiling.html">Profiling</a></li> </ul> </details></li> <li class="toctree-l2 current active"><a class="current reference internal" href="#">flax.serialization package</a></li> <li class="toctree-l2"><a class="reference internal" href="flax.struct.html">flax.struct package</a></li> <li class="toctree-l2"><a class="reference internal" href="flax.traceback_util.html">flax.traceback_util package</a></li> <li class="toctree-l2"><a class="reference internal" href="flax.training.html">flax.training package</a></li> <li class="toctree-l2"><a class="reference internal" href="flax.traverse_util.html">flax.traverse_util package</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference external" href="https://flax.readthedocs.io/en/latest/index.html">Flax NNX</a></li> </ul> </div> </nav></div> </div> <div class="sidebar-primary-items__end sidebar-primary__section"> </div> <div id="rtd-footer-container"></div> </div> <main id="main-content" class="bd-main" role="main"> <div class="sbt-scroll-pixel-helper"></div> <div class="bd-content"> <div class="bd-article-container"> <div class="bd-header-article d-print-none"> <div class="header-article-items header-article__inner"> <div class="header-article-items__start"> <div class="header-article-item"><button class="sidebar-toggle primary-toggle btn btn-sm" title="Toggle primary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip"> <span class="fa-solid fa-bars"></span> </button></div> </div> <div class="header-article-items__end"> <div class="header-article-item"> <div class="article-header-buttons"> <a href="https://github.com/google/flax" target="_blank" class="btn btn-sm btn-source-repository-button" title="Source repository" data-bs-placement="bottom" data-bs-toggle="tooltip" > <span class="btn__icon-container"> <i class="fab fa-github"></i> </span> </a> <div class="dropdown dropdown-download-buttons"> <button class="btn dropdown-toggle" type="button" data-bs-toggle="dropdown" aria-expanded="false" aria-label="Download this page"> <i class="fas fa-download"></i> </button> <ul class="dropdown-menu"> <li><a href="../_sources/api_reference/flax.serialization.rst" target="_blank" class="btn btn-sm btn-download-source-button dropdown-item" title="Download source file" data-bs-placement="left" data-bs-toggle="tooltip" > <span class="btn__icon-container"> <i class="fas fa-file"></i> </span> <span class="btn__text-container">.rst</span> </a> </li> <li> <button onclick="window.print()" class="btn btn-sm btn-download-pdf-button dropdown-item" title="Print to PDF" data-bs-placement="left" data-bs-toggle="tooltip" > <span class="btn__icon-container"> <i class="fas fa-file-pdf"></i> </span> <span class="btn__text-container">.pdf</span> </button> </li> </ul> </div> <button onclick="toggleFullScreen()" class="btn btn-sm btn-fullscreen-button" title="Fullscreen mode" data-bs-placement="bottom" data-bs-toggle="tooltip" > <span class="btn__icon-container"> <i class="fas fa-expand"></i> </span> </button> <script> document.write(` <button class="btn btn-sm nav-link pst-navbar-icon theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip"> <i class="theme-switch fa-solid fa-sun fa-lg" data-mode="light"></i> <i class="theme-switch fa-solid fa-moon fa-lg" data-mode="dark"></i> <i class="theme-switch fa-solid fa-circle-half-stroke fa-lg" data-mode="auto"></i> </button> `); </script> <script> document.write(` <button class="btn btn-sm pst-navbar-icon search-button search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip"> <i class="fa-solid fa-magnifying-glass fa-lg"></i> </button> `); </script> <button class="sidebar-toggle secondary-toggle btn btn-sm" title="Toggle secondary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip"> <span class="fa-solid fa-list"></span> </button> </div></div> </div> </div> </div> <div id="jb-print-docs-body" class="onlyprint"> <h1>flax.serialization package</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="#state-dicts">State dicts</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.from_state_dict"><code class="docutils literal notranslate"><span class="pre">from_state_dict()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.to_state_dict"><code class="docutils literal notranslate"><span class="pre">to_state_dict()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.register_serialization_state"><code class="docutils literal notranslate"><span class="pre">register_serialization_state()</span></code></a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#serialization-with-messagepack">Serialization with MessagePack</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.msgpack_serialize"><code class="docutils literal notranslate"><span class="pre">msgpack_serialize()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.msgpack_restore"><code class="docutils literal notranslate"><span class="pre">msgpack_restore()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.to_bytes"><code class="docutils literal notranslate"><span class="pre">to_bytes()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.from_bytes"><code class="docutils literal notranslate"><span class="pre">from_bytes()</span></code></a></li> </ul> </li> </ul> </nav> </div> </div> </div> <div id="searchbox"></div> <article class="bd-article"> <div class="section" id="module-flax.serialization"> <span id="flax-serialization-package"></span><h1>flax.serialization package<a class="headerlink" href="#module-flax.serialization" title="Permalink to this heading">#</a></h1> <p>Serialization utilities for Jax.</p> <p>All Flax classes that carry state (e.g., Optimizer) can be turned into a state dict of numpy arrays for easy serialization.</p> <div class="section" id="state-dicts"> <h2>State dicts<a class="headerlink" href="#state-dicts" title="Permalink to this heading">#</a></h2> <dl class="py function"> <dt class="sig sig-object py" id="flax.serialization.from_state_dict"> <span class="sig-prename descclassname"><span class="pre">flax.serialization.</span></span><span class="sig-name descname"><span class="pre">from_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">target</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">state</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">name</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">'.'</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/flax/serialization.html#from_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#flax.serialization.from_state_dict" title="Permalink to this definition">#</a></dt> <dd><p>Restores the state of the given target using a state dict.</p> <p>This function takes the current target as an argument. This lets us know the exact structure of the target, as well as lets us add assertions that shapes and dtypes don’t change.</p> <p>In practice, none of the leaf values in <code class="docutils literal notranslate"><span class="pre">target</span></code> are actually used. Only the tree structure, shapes and dtypes.</p> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li><p><strong>target</strong> – the object of which the state should be restored.</p></li> <li><p><strong>state</strong> – a dictionary generated by <code class="docutils literal notranslate"><span class="pre">to_state_dict</span></code> with the desired new state for <code class="docutils literal notranslate"><span class="pre">target</span></code>.</p></li> <li><p><strong>name</strong> – name of branch taken, used to improve deserialization error messages.</p></li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even"><p>A copy of the object with the restored state.</p> </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.serialization.to_state_dict"> <span class="sig-prename descclassname"><span class="pre">flax.serialization.</span></span><span class="sig-name descname"><span class="pre">to_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">target</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/flax/serialization.html#to_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#flax.serialization.to_state_dict" title="Permalink to this definition">#</a></dt> <dd><p>Returns a dictionary with the state of the given target.</p> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.serialization.register_serialization_state"> <span class="sig-prename descclassname"><span class="pre">flax.serialization.</span></span><span class="sig-name descname"><span class="pre">register_serialization_state</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">ty</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">ty_to_state_dict</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">ty_from_state_dict</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">override</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/flax/serialization.html#register_serialization_state"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#flax.serialization.register_serialization_state" title="Permalink to this definition">#</a></dt> <dd><p>Register a type for serialization.</p> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li><p><strong>ty</strong> – the type to be registered</p></li> <li><p><strong>ty_to_state_dict</strong> – a function that takes an instance of ty and returns its state as a dictionary.</p></li> <li><p><strong>ty_from_state_dict</strong> – a function that takes an instance of ty and a state dict, and returns a copy of the instance with the restored state.</p></li> <li><p><strong>override</strong> – override a previously registered serialization handler (default: False).</p></li> </ul> </dd> </dl> </dd></dl> </div> <div class="section" id="serialization-with-messagepack"> <h2>Serialization with MessagePack<a class="headerlink" href="#serialization-with-messagepack" title="Permalink to this heading">#</a></h2> <dl class="py function"> <dt class="sig sig-object py" id="flax.serialization.msgpack_serialize"> <span class="sig-prename descclassname"><span class="pre">flax.serialization.</span></span><span class="sig-name descname"><span class="pre">msgpack_serialize</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">pytree</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">in_place</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/flax/serialization.html#msgpack_serialize"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#flax.serialization.msgpack_serialize" title="Permalink to this definition">#</a></dt> <dd><p>Save data structure to bytes in msgpack format.</p> <p>Low-level function that only supports python trees with array leaves, for custom objects use <code class="docutils literal notranslate"><span class="pre">to_bytes</span></code>. It splits arrays above MAX_CHUNK_SIZE into multiple chunks.</p> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li><p><strong>pytree</strong> – python tree of dict, list, tuple with python primitives and array leaves.</p></li> <li><p><strong>in_place</strong> – boolean specifying if pytree should be modified in place.</p></li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even"><p>msgpack-encoded bytes of pytree.</p> </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.serialization.msgpack_restore"> <span class="sig-prename descclassname"><span class="pre">flax.serialization.</span></span><span class="sig-name descname"><span class="pre">msgpack_restore</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">encoded_pytree</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/flax/serialization.html#msgpack_restore"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#flax.serialization.msgpack_restore" title="Permalink to this definition">#</a></dt> <dd><p>Restore data structure from bytes in msgpack format.</p> <p>Low-level function that only supports python trees with array leaves, for custom objects use <code class="docutils literal notranslate"><span class="pre">from_bytes</span></code>.</p> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><p><strong>encoded_pytree</strong> – msgpack-encoded bytes of python tree.</p> </dd> <dt class="field-even">Returns</dt> <dd class="field-even"><p>Python tree of dict, list, tuple with python primitive and array leaves.</p> </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.serialization.to_bytes"> <span class="sig-prename descclassname"><span class="pre">flax.serialization.</span></span><span class="sig-name descname"><span class="pre">to_bytes</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">target</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/flax/serialization.html#to_bytes"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#flax.serialization.to_bytes" title="Permalink to this definition">#</a></dt> <dd><p>Save optimizer or other object as msgpack-serialized state-dict.</p> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><p><strong>target</strong> – template object with state-dict registrations to be serialized to msgpack format. Typically a flax model or optimizer.</p> </dd> <dt class="field-even">Returns</dt> <dd class="field-even"><p>Bytes of msgpack-encoded state-dict of <code class="docutils literal notranslate"><span class="pre">target</span></code> object.</p> </dd> </dl> </dd></dl> <dl class="py function"> <dt class="sig sig-object py" id="flax.serialization.from_bytes"> <span class="sig-prename descclassname"><span class="pre">flax.serialization.</span></span><span class="sig-name descname"><span class="pre">from_bytes</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">target</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">encoded_bytes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/flax/serialization.html#from_bytes"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#flax.serialization.from_bytes" title="Permalink to this definition">#</a></dt> <dd><p>Restore optimizer or other object from msgpack-serialized state-dict.</p> <dl class="field-list simple"> <dt class="field-odd">Parameters</dt> <dd class="field-odd"><ul class="simple"> <li><p><strong>target</strong> – template object with state-dict registrations that matches the structure being deserialized from <code class="docutils literal notranslate"><span class="pre">encoded_bytes</span></code>.</p></li> <li><p><strong>encoded_bytes</strong> – msgpack serialized object structurally isomorphic to <code class="docutils literal notranslate"><span class="pre">target</span></code>. Typically a flax model or optimizer.</p></li> </ul> </dd> <dt class="field-even">Returns</dt> <dd class="field-even"><p>A new object structurally isomorphic to <code class="docutils literal notranslate"><span class="pre">target</span></code> containing the updated leaf data from saved data.</p> </dd> </dl> </dd></dl> </div> </div> </article> <footer class="prev-next-footer d-print-none"> <div class="prev-next-area"> <a class="left-prev" href="flax.linen/profiling.html" title="previous page"> <i class="fa-solid fa-angle-left"></i> <div class="prev-next-info"> <p class="prev-next-subtitle">previous</p> <p class="prev-next-title">Profiling</p> </div> </a> <a class="right-next" href="flax.struct.html" title="next page"> <div class="prev-next-info"> <p class="prev-next-subtitle">next</p> <p class="prev-next-title">flax.struct package</p> </div> <i class="fa-solid fa-angle-right"></i> </a> </div> </footer> </div> <div class="bd-sidebar-secondary bd-toc"><div class="sidebar-secondary-items sidebar-secondary__inner"> <div class="sidebar-secondary-item"> <div class="page-toc tocsection onthispage"> <i class="fa-solid fa-list"></i> Contents </div> <nav class="bd-toc-nav page-toc"> <ul class="visible nav section-nav flex-column"> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#state-dicts">State dicts</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.from_state_dict"><code class="docutils literal notranslate"><span class="pre">from_state_dict()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.to_state_dict"><code class="docutils literal notranslate"><span class="pre">to_state_dict()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.register_serialization_state"><code class="docutils literal notranslate"><span class="pre">register_serialization_state()</span></code></a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#serialization-with-messagepack">Serialization with MessagePack</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.msgpack_serialize"><code class="docutils literal notranslate"><span class="pre">msgpack_serialize()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.msgpack_restore"><code class="docutils literal notranslate"><span class="pre">msgpack_restore()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.to_bytes"><code class="docutils literal notranslate"><span class="pre">to_bytes()</span></code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#flax.serialization.from_bytes"><code class="docutils literal notranslate"><span class="pre">from_bytes()</span></code></a></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"> <p class="component-author"> By The Flax authors </p> </div> <div class="footer-item"> <p class="copyright"> © Copyright 2023, The Flax authors. <br/> </p> </div> <div class="footer-item"> </div> <div class="footer-item"> </div> </div> </footer> </main> </div> </div>  <script src="../_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b"></script> <script src="../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b"></script> <footer class="bd-footer"> </footer> </body> </html>