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>How to contribute</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 = 'contributing';</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="API Reference" href="api_reference/index.html" /> <link rel="prev" title="The Flax philosophy" href="philosophy.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="/contributing.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 current active"><a class="current reference internal" href="#">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><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="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><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="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"> <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/contributing.md" 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">.md</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>How to contribute</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="#ways-to-contribute">Ways to contribute</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#contributing-code-using-pull-requests">Contributing code using pull requests</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#to-create-a-pull-request-from-a-fork">To create a pull request from a fork</a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#adding-or-updating-dependencies">Adding or updating dependencies</a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#updating-jupyter-notebooks">Updating Jupyter Notebooks</a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#editing-jupyter-notebooks-ipynb">Editing Jupyter Notebooks (<code class="docutils literal notranslate"><span class="pre">.ipynb</span></code>)</a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#editing-markdown-files-md">Editing Markdown files (<code class="docutils literal notranslate"><span class="pre">.md</span></code>)</a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#syncing-notebooks">Syncing notebooks</a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#creating-new-notebooks">Creating new notebooks</a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#notebooks-within-the-sphinx-build">Notebooks within the Sphinx build</a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#updating-the-pull-request-contents">Updating the pull request contents</a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#troubleshooting">Troubleshooting</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#too-many-commits-in-a-pull-request">Too many commits in a pull request</a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#contributor-license-agreement">Contributor License Agreement</a></li> </ul> </nav> </div> </div> </div> <div id="searchbox"></div> <article class="bd-article"> <div class="tex2jax_ignore mathjax_ignore section" id="how-to-contribute"> <h1>How to contribute<a class="headerlink" href="#how-to-contribute" title="Permalink to this heading">#</a></h1> <p>Everyone can contribute to Flax, and the Flax development team values everyone’s contributions! You can contribute in many more ways than just writing code. Answering questions on the <a class="reference external" href="https://github.com/google/flax/discussions">Flax GitHub Discussions page</a>, helping each other, and improving Flax documentation are extremely valuable to the Flax ecosystem.</p> <p>We also appreciate if you spread the word, for instance by starring the <a class="reference external" href="https://github.com/google/flax">Flax GitHub repository</a>, or referencing Flax in blog posts of projects that used it.</p> <p>This project follows <a class="reference external" href="https://opensource.google/conduct/">Google’s Open Source Community Guidelines</a>.</p> <div class="section" id="ways-to-contribute"> <h2>Ways to contribute<a class="headerlink" href="#ways-to-contribute" title="Permalink to this heading">#</a></h2> <p>We welcome pull requests (PRs), in particular for those issues <a class="reference external" href="https://github.com/google/flax/issues?q=is%3Aopen+is%3Aissue+label%3A%22Status%3A+pull+requests+welcome%22">marked as PR-ready</a>. For other proposals, you should first open a GitHub Issue or a GitHub Discussion to start a conversation about your planned contribution.</p> </div> <div class="section" id="contributing-code-using-pull-requests"> <h2>Contributing code using pull requests<a class="headerlink" href="#contributing-code-using-pull-requests" title="Permalink to this heading">#</a></h2> <p>The Flax development team performs all development using <a class="reference external" href="https://git-scm.com/">Git</a>. To contribute, you should have basic knowledge of <a class="reference external" href="https://git-scm.com/">Git</a> and <a class="reference external" href="https://docs.github.com">GitHub</a>. (You can learn how to set up Git by following Git’s official <a class="reference external" href="https://git-scm.com/book/en/v2/Getting-Started-First-Time-Git-Setup">Getting Started - First-Time Git Setup</a> and GitHub’s <a class="reference external" href="https://docs.github.com/en/get-started/quickstart/set-up-git">Set Up Git</a> guides.)</p> <p>To contribute code to Flax on GitHub, follow these steps:</p> <div class="section" id="to-create-a-pull-request-from-a-fork"> <h3>To create a pull request from a fork<a class="headerlink" href="#to-create-a-pull-request-from-a-fork" title="Permalink to this heading">#</a></h3> <ol class="arabic"> <li><p>Using GitHub’s web UI, fork the Flax repository by clicking the ‘Fork’ button on the <a class="reference external" href="http://www.github.com/google/flax"><code class="docutils literal notranslate"><span class="pre">github.com/google/flax</span></code> repository page</a>. This creates a fork (a copy) of the Flax repository in your own GitHub.</p> <p>Reference: <a class="reference external" href="https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request-from-a-fork">Creating a pull request from a fork</a>.</p> </li> <li><p>Install <a class="reference external" href="https://www.python.org/downloads/">Python >=3.7</a>.</p></li> <li><p>(Optional) Create a virtual environment or a Docker container. See <a class="reference external" href="https://github.com/google/flax/blob/main/dev/README.md"><code class="docutils literal notranslate"><span class="pre">dev/README.md</span></code></a> for details on how to set up a Docker Container. To set up a virtual environment, run the following:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>python3<span class="w"> </span>-m<span class="w"> </span>virtualenv<span class="w"> </span>env .<span class="w"> </span>env/bin/activate </pre></div> </div> <p>This ensures all your dependencies are installed in this environment.</p> </li> <li><p>Clone your local forked Flax repo with <code class="docutils literal notranslate"><span class="pre">git</span> <span class="pre">clone</span></code>. Then, install the required packages with <a class="reference external" href="https://pip.pypa.io/en/stable/cli/pip_install/">PyPi</a>. This enables you to immediately test the code after modifying it:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>git<span class="w"> </span>clone<span class="w"> </span>https://github.com/YOUR_USERNAME/flax <span class="nb">cd</span><span class="w"> </span>flax pip<span class="w"> </span>install<span class="w"> </span>-e<span class="w"> </span><span class="s2">".[all,testing,docs]"</span> </pre></div> </div> <p>You can also use <a class="reference external" href="https://docs.astral.sh/uv/">uv</a> to setup the development environment:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>uv<span class="w"> </span>sync<span class="w"> </span>--all-extras </pre></div> </div> </li> <li><p>Set up pre-commit hooks, this will run some automated checks during each <code class="docutils literal notranslate"><span class="pre">git</span></code> commit and possibly update some files that require changes.</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>pip<span class="w"> </span>install<span class="w"> </span>pre-commit pre-commit<span class="w"> </span>install </pre></div> </div> </li> <li><p>Add the Google Flax repo (not your fork) as an upstream remote, so you can use it to sync your changes.</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>git<span class="w"> </span>remote<span class="w"> </span>add<span class="w"> </span>upstream<span class="w"> </span>http://www.github.com/google/flax </pre></div> </div> </li> <li><p>Create a branch, such as <code class="docutils literal notranslate"><span class="pre">my_development_branch</span></code>, you will develop from:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>git<span class="w"> </span>checkout<span class="w"> </span>-b<span class="w"> </span>my_development_branch </pre></div> </div> </li> <li><p>Implement your changes using your favorite editor (we recommend <a class="reference external" href="https://code.visualstudio.com/">Visual Studio Code</a>).</p> <p>Make sure the tests pass by running the following command from the top of the repository:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>./tests/run_all_tests.sh </pre></div> </div> </li> <li><p>Once you finish making changes, don’t forget to create commits (<a class="reference external" href="https://chris.beams.io/posts/git-commit/">learn how to write a commit message</a>):</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>git<span class="w"> </span>add<span class="w"> </span>file1.py<span class="w"> </span>file2.py<span class="w"> </span>... <span class="c1"># or use `git add .` to add all changed files</span> git<span class="w"> </span>commit<span class="w"> </span>-m<span class="w"> </span><span class="s2">"Your commit message"</span> </pre></div> </div> <p>Then sync your code with the main repository:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>git<span class="w"> </span>fetch<span class="w"> </span>upstream git<span class="w"> </span>rebase<span class="w"> </span>upstream/main </pre></div> </div> </li> <li><p>Finally, push your commit on your <code class="docutils literal notranslate"><span class="pre">my_development_branch</span></code>, and create a remote branch in your fork that you can use to create a pull request from:</p></li> </ol> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>git<span class="w"> </span>push<span class="w"> </span>--set-upstream<span class="w"> </span>origin<span class="w"> </span>my_development_branch </pre></div> </div> <p>After running the command, you should get a GitHub link in your (VS Code) terminal output for creating a pull request. If you don’t receive a link after <code class="docutils literal notranslate"><span class="pre">git</span> <span class="pre">push</span></code>, use the <a class="reference external" href="https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request?tool=webui">GitHub web UI</a> to create a pull request.</p> <ol class="arabic simple" start="11"> <li><p>Make sure your pull request passes the <a class="reference external" href="https://github.com/google/flax/blob/main/.github/pull_request_template.md#checklist">Flax PR checklist</a>. If so, create a pull request from the Flax repository and send it for review. Consult <a class="reference external" href="https://help.github.com/articles/about-pull-requests/">GitHub Help</a> for more information on using pull requests.</p></li> </ol> <p>You can learn more in GitHub’s <a class="reference external" href="https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request-from-a-fork">Creating a pull request from a fork </a>. documentation.</p> </div> <div class="section" id="adding-or-updating-dependencies"> <h3>Adding or updating dependencies<a class="headerlink" href="#adding-or-updating-dependencies" title="Permalink to this heading">#</a></h3> <p>To add or update dependencies, you must use <code class="docutils literal notranslate"><span class="pre">uv</span></code> after updating the <code class="docutils literal notranslate"><span class="pre">pyproject.toml</span></code> file to ensure that the <code class="docutils literal notranslate"><span class="pre">uv.lock</span></code> file is up-to-date.</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>uv<span class="w"> </span>sync<span class="w"> </span>--all-extras </pre></div> </div> <p>Alternatively use can use <code class="docutils literal notranslate"><span class="pre">uv</span> <span class="pre">add</span></code> to add or update the dependencies automatically, for example:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>uv<span class="w"> </span>add<span class="w"> </span><span class="s1">'some-package>=1.2.3'</span> </pre></div> </div> </div> <div class="section" id="updating-jupyter-notebooks"> <h3>Updating Jupyter Notebooks<a class="headerlink" href="#updating-jupyter-notebooks" title="Permalink to this heading">#</a></h3> <p>We use <a class="reference external" href="https://jupytext.readthedocs.io/">jupytext</a> to maintain two synced copies of docs in <code class="docutils literal notranslate"><span class="pre">docs/notebooks</span></code>: one in the Jupyter Notebook (<code class="docutils literal notranslate"><span class="pre">.ipynb</span></code>) format, and one in Markdown (<code class="docutils literal notranslate"><span class="pre">.md</span></code>).</p> <p>The former can be opened and executed directly in <a class="reference external" href="https://colab.research.google.com/">Google Colab</a>. Markdown makes it easier to track changes/diffs within version control and, for example, GitHub web UI, since <code class="docutils literal notranslate"><span class="pre">.ipynb</span></code> files are based on JSON.</p> <div class="section" id="editing-jupyter-notebooks-ipynb"> <h4>Editing Jupyter Notebooks (<code class="docutils literal notranslate"><span class="pre">.ipynb</span></code>)<a class="headerlink" href="#editing-jupyter-notebooks-ipynb" title="Permalink to this heading">#</a></h4> <p>For making large changes that substantially modify code and outputs, it’s recommended to edit the notebooks in <a class="reference external" href="https://jupyter.org/install">Jupyter</a> or in <a class="reference external" href="https://colab.research.google.com/">Colab</a>.</p> <p>If you choose to work in Colab, go to <strong>File</strong> and click <strong>Upload notebook</strong>, then pick your file. After loading it into Colab and editing it, make sure you run the cells, and that there aren’t any errors. Click on <strong>Runtime</strong>, then select <strong>Run all</strong>. After you finish, click <strong>File</strong> > <strong>Download</strong> > <strong>Download ipynb</strong>. You may also want to test that the file executes properly by using <code class="docutils literal notranslate"><span class="pre">sphinx-build</span></code>, as explained above.</p> <p>After you make changes in your Jupyter Notebook, follow the steps <em>Syncing notebooks</em> below.</p> </div> <div class="section" id="editing-markdown-files-md"> <h4>Editing Markdown files (<code class="docutils literal notranslate"><span class="pre">.md</span></code>)<a class="headerlink" href="#editing-markdown-files-md" title="Permalink to this heading">#</a></h4> <p>For making smaller changes to the text content of the notebooks, it is easiest to edit the <code class="docutils literal notranslate"><span class="pre">.md</span></code> versions using a text editor.</p> <p>After you make changes in your Markdown file, follow the steps <em>Syncing notebooks</em> below.</p> </div> <div class="section" id="syncing-notebooks"> <h4>Syncing notebooks<a class="headerlink" href="#syncing-notebooks" title="Permalink to this heading">#</a></h4> <p>After editing either the <code class="docutils literal notranslate"><span class="pre">.ipynb</span></code> or <code class="docutils literal notranslate"><span class="pre">.md</span></code> versions of the docs, sync the two versions using <a class="reference external" href="https://jupytext.readthedocs.io/">jupytext</a> by running <code class="docutils literal notranslate"><span class="pre">jupytext</span> <span class="pre">--sync</span></code> on the updated notebooks.</p> <p>First, make sure you have jupytext installed. The jupytext version should match the one specified in <a class="reference external" href="https://github.com/google/flax/blob/main/.pre-commit-config.yaml">.pre-commit-config.yaml</a> (currently, it is v1.13.8).</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>pip<span class="w"> </span>install<span class="w"> </span><span class="nv">jupytext</span><span class="o">==</span><span class="m">1</span>.13.8 </pre></div> </div> <p>Then, after you have made your changes in the Jupyter Notebook, sync the contents with its Markdown-equivalent file by running the following command:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>jupytext<span class="w"> </span>--sync<span class="w"> </span>path/to/the/file.ipynb </pre></div> </div> <p>Similarly, to sync your Markdown file with its Jupyter Notebook version, run:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>jupytext<span class="w"> </span>--sync<span class="w"> </span>path/to/the/file.md </pre></div> </div> <p>Note that if you receive an error, and it is the first time you worked in a Jupyter Notebook, you may need to (re)create a synced copy of the document (which is explained in detail in <em>Creating new notebooks</em> section below):</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>jupytext<span class="w"> </span>--set-formats<span class="w"> </span>ipynb,md:myst<span class="w"> </span>path/to/the/notebook.ipynb </pre></div> </div> <p>Once you’re finished with syncing the <code class="docutils literal notranslate"><span class="pre">.md</span></code> and <code class="docutils literal notranslate"><span class="pre">.ipynb</span></code> files, you can check that they are properly synced using the <a class="reference external" href="https://pre-commit.com/">pre-commit</a> framework to perform the same checks used in the Flax GitHub CI:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>git<span class="w"> </span>add<span class="w"> </span>docs<span class="w"> </span>-u<span class="w"> </span><span class="c1"># pre-commit runs on files in git staging.</span> pre-commit<span class="w"> </span>run<span class="w"> </span>jupytext </pre></div> </div> </div> <div class="section" id="creating-new-notebooks"> <h4>Creating new notebooks<a class="headerlink" href="#creating-new-notebooks" title="Permalink to this heading">#</a></h4> <p>If you are adding a new Jupyter Notebook to the documentation, you can use <code class="docutils literal notranslate"><span class="pre">jupytext</span> <span class="pre">--set-formats</span></code>. It can set up both the Jupyter Notebook (<code class="docutils literal notranslate"><span class="pre">.ipynb</span></code>) and Markdown (<code class="docutils literal notranslate"><span class="pre">.md</span></code>) versions of the file:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>jupytext<span class="w"> </span>--set-formats<span class="w"> </span>ipynb,md:myst<span class="w"> </span>path/to/the/notebook.ipynb </pre></div> </div> <p>This works by adding a <code class="docutils literal notranslate"><span class="pre">"jupytext"</span></code> metadata field to the notebook file which specifies the desired formats. The <code class="docutils literal notranslate"><span class="pre">jupytext</span> <span class="pre">--sync</span></code> command can then recognize them when invoked.</p> <p>After you make changes in your file(s), follow the steps from the <em>Syncing notebooks</em> section above to keep the contents of both Markdown and Jupyter Notebook files in sync.</p> </div> <div class="section" id="notebooks-within-the-sphinx-build"> <h4>Notebooks within the Sphinx build<a class="headerlink" href="#notebooks-within-the-sphinx-build" title="Permalink to this heading">#</a></h4> <p>Some of the notebooks are built automatically as part of the pre-submit checks and as part of the <a class="reference external" href="https://flax.readthedocs.io/en/latest">Read the Docs</a> build. The build will fail if cells raise errors. If the errors are intentional, you can either catch them, or tag the cell with <code class="docutils literal notranslate"><span class="pre">raises-exceptions</span></code> metadata (<a class="reference external" href="https://github.com/jax-ml/jax/pull/2402/files">example PR</a>). You have to add this metadata by hand in the <code class="docutils literal notranslate"><span class="pre">.ipynb</span></code> file. It will be preserved when somebody else re-saves the notebook.</p> <p>We exclude some notebooks from the build because, for example, they contain long computations. See <code class="docutils literal notranslate"><span class="pre">exclude_patterns</span></code> in <a class="reference external" href="https://github.com/google/flax/blob/main/docs/conf.py"><code class="docutils literal notranslate"><span class="pre">conf.py</span></code></a>.</p> </div> </div> <div class="section" id="updating-the-pull-request-contents"> <h3>Updating the pull request contents<a class="headerlink" href="#updating-the-pull-request-contents" title="Permalink to this heading">#</a></h3> <p>Every pull request should ideally be limited to just one commit, so if you have multiple commits please squash them.</p> <p>Assuming you now have only one commit in your pull request, and want to add changes requested during review:</p> <ol class="arabic simple"> <li><p>Make the changes locally in your editor.</p></li> <li><p>Run <code class="docutils literal notranslate"><span class="pre">git</span> <span class="pre">commit</span> <span class="pre">-a</span> <span class="pre">--amend</span></code>. This updates the commit contents and allows you to edit the commit message.</p></li> <li><p>At this point, <code class="docutils literal notranslate"><span class="pre">git</span> <span class="pre">push</span></code> alone will result in an error. Instead, use <code class="docutils literal notranslate"><span class="pre">git</span> <span class="pre">push</span> <span class="pre">--force</span></code>.</p></li> <li><p>Check that it’s done: The changes to your commit should be immediately reflected in the Github web UI.</p></li> </ol> </div> </div> <div class="section" id="troubleshooting"> <h2>Troubleshooting<a class="headerlink" href="#troubleshooting" title="Permalink to this heading">#</a></h2> <div class="section" id="too-many-commits-in-a-pull-request"> <h3>Too many commits in a pull request<a class="headerlink" href="#too-many-commits-in-a-pull-request" title="Permalink to this heading">#</a></h3> <p>If your PR has too many commits associated with it (for example, more than five), you need to squash them. Otherwise, the Flax docs build process may fail with an error message. This is because of the following reasons:</p> <ul class="simple"> <li><p>There are more than five commits in your pull request; and</p></li> <li><p>The Flax source sync process fails when the commit tree is too large.</p></li> </ul> <p>To squash your commits, you can rebase your branch to <code class="docutils literal notranslate"><span class="pre">main</span></code> and create a new commit containing all your changes, run the following command:</p> <div class="highlight-bash notranslate"><div class="highlight"><pre><span></span>git<span class="w"> </span>rebase<span class="w"> </span>main<span class="w"> </span><span class="o">&&</span><span class="w"> </span>git<span class="w"> </span>reset<span class="w"> </span>--soft<span class="w"> </span>main<span class="w"> </span><span class="o">&&</span><span class="w"> </span>git<span class="w"> </span>commit </pre></div> </div> <p>This will apply all your changes to the main branch. Note that if you had to resolve any conflicts while working on your change (for instance, you did a <code class="docutils literal notranslate"><span class="pre">pull</span> <span class="pre">upstream</span> <span class="pre">main</span></code> which led to conflict), then you will have to resolve these conflicts again.</p> <p>After you have successfully rebased your branch, you should push your changes. And because you changed the commit history, you may have to use <code class="docutils literal notranslate"><span class="pre">git</span> <span class="pre">push</span> <span class="pre">--force</span></code>.</p> </div> </div> <div class="section" id="contributor-license-agreement"> <h2>Contributor License Agreement<a class="headerlink" href="#contributor-license-agreement" title="Permalink to this heading">#</a></h2> <p>Contributions to this project must be accompanied by a Contributor License Agreement. You (or your employer) retain the copyright to your contribution; this simply gives us permission to use and redistribute your contributions as part of the project. Head over to <a class="reference external" href="https://cla.developers.google.com/">https://cla.developers.google.com/</a> to see your current agreements on file or to sign a new one.</p> <p>You generally only need to submit a CLA once, so if you’ve already submitted one (even if it was for a different project), you probably don’t need to do it again.</p> </div> </div> </article> <footer class="prev-next-footer d-print-none"> <div class="prev-next-area"> <a class="left-prev" href="philosophy.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">The Flax philosophy</p> </div> </a> <a class="right-next" href="api_reference/index.html" title="next page"> <div class="prev-next-info"> <p class="prev-next-subtitle">next</p> <p class="prev-next-title">API Reference</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="#ways-to-contribute">Ways to contribute</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#contributing-code-using-pull-requests">Contributing code using pull requests</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#to-create-a-pull-request-from-a-fork">To create a pull request from a fork</a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#adding-or-updating-dependencies">Adding or updating dependencies</a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#updating-jupyter-notebooks">Updating Jupyter Notebooks</a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#editing-jupyter-notebooks-ipynb">Editing Jupyter Notebooks (<code class="docutils literal notranslate"><span class="pre">.ipynb</span></code>)</a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#editing-markdown-files-md">Editing Markdown files (<code class="docutils literal notranslate"><span class="pre">.md</span></code>)</a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#syncing-notebooks">Syncing notebooks</a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#creating-new-notebooks">Creating new notebooks</a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#notebooks-within-the-sphinx-build">Notebooks within the Sphinx build</a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#updating-the-pull-request-contents">Updating the pull request contents</a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#troubleshooting">Troubleshooting</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#too-many-commits-in-a-pull-request">Too many commits in a pull request</a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#contributor-license-agreement">Contributor License Agreement</a></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>