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" /><meta name="viewport" content="width=device-width, initial-scale=1" /> <title>jax.sharding module — JAX documentation</title> <script data-cfasync="false"> document.documentElement.dataset.mode = localStorage.getItem("mode") || ""; document.documentElement.dataset.theme = localStorage.getItem("theme") || "light"; </script>  <link href="_static/styles/theme.css?digest=5b4479735964841361fd" rel="stylesheet" /> <link href="_static/styles/bootstrap.css?digest=5b4479735964841361fd" rel="stylesheet" /> <link href="_static/styles/pydata-sphinx-theme.css?digest=5b4479735964841361fd" rel="stylesheet" /> <link href="_static/vendor/fontawesome/6.1.2/css/all.min.css?digest=5b4479735964841361fd" rel="stylesheet" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="_static/vendor/fontawesome/6.1.2/webfonts/fa-solid-900.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="_static/vendor/fontawesome/6.1.2/webfonts/fa-brands-400.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="_static/vendor/fontawesome/6.1.2/webfonts/fa-regular-400.woff2" /> <link rel="stylesheet" type="text/css" href="_static/pygments.css?v=8f2a1f02" /> <link rel="stylesheet" type="text/css" href="_static/styles/sphinx-book-theme.css?v=384b581d" /> <link rel="stylesheet" type="text/css" href="_static/plot_directive.css" /> <link rel="stylesheet" type="text/css" href="_static/mystnb.4510f1fc1dee50b3e5859aac5469c37c29e427902b24a333a5f9fcb2f0b3ac41.css" /> <link rel="stylesheet" type="text/css" href="_static/copybutton.css?v=76b2166b" /> <link rel="stylesheet" type="text/css" href="_static/sphinx-design.min.css?v=95c83b7e" /> <link rel="stylesheet" type="text/css" href="_static/style.css?v=05b22d1f" /> <link rel="stylesheet" href="_static/style.css" type="text/css" />  <link rel="preload" as="script" href="_static/scripts/bootstrap.js?digest=5b4479735964841361fd" /> <link rel="preload" as="script" href="_static/scripts/pydata-sphinx-theme.js?digest=5b4479735964841361fd" /> <script src="_static/vendor/fontawesome/6.1.2/js/all.min.js?digest=5b4479735964841361fd"></script> <script src="_static/documentation_options.js?v=9eb32ce0"></script> <script src="_static/doctools.js?v=9a2dae69"></script> <script src="_static/sphinx_highlight.js?v=dc90522c"></script> <script src="_static/clipboard.min.js?v=a7894cd8"></script> <script src="_static/copybutton.js?v=30646c52"></script> <script src="_static/scripts/sphinx-book-theme.js?v=efea14e4"></script> <script src="_static/design-tabs.js?v=f930bc37"></script> <script>DOCUMENTATION_OPTIONS.pagename = 'jax.sharding';</script> <link rel="icon" href="_static/favicon.png"/> <link rel="author" title="About these documents" href="about.html" /> <link rel="index" title="Index" href="genindex.html" /> <link rel="search" title="Search" href="search.html" /> <link rel="next" title="jax.debug module" href="jax.debug.html" /> <link rel="prev" title="jax.random.weibull_min" href="_autosummary/jax.random.weibull_min.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="jax" /><meta name="readthedocs-version-slug" content="latest" /><meta name="readthedocs-resolver-filename" content="/jax.sharding.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=""> <a class="skip-link" href="#main-content">Skip to main content</a> <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" name="__primary" id="__primary"/> <label class="overlay overlay-primary" for="__primary"></label> <input type="checkbox" class="sidebar-toggle" name="__secondary" id="__secondary"/> <label class="overlay overlay-secondary" for="__secondary"></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> <nav class="bd-header navbar navbar-expand-lg bd-navbar"> </nav> <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/jax_logo_250px.png" class="logo__image only-light" alt="JAX documentation - Home"/> <script>document.write(`<img src="_static/jax_logo_250px.png" class="logo__image only-dark" alt="JAX documentation - Home"/>`);</script> </a></div> <div class="sidebar-primary-item"> <script> document.write(` <button class="btn navbar-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"> Getting started <ul class="nav bd-sidenav"> <li class="toctree-l1"><a class="reference internal" href="installation.html">Installation</a></li> <li class="toctree-l1"><a class="reference internal" href="quickstart.html">Quickstart</a></li> </ul> <ul class="nav bd-sidenav"> <li class="toctree-l1 has-children"><a class="reference internal" href="tutorials.html">Tutorials</a><input class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-1"></label><ul> <li class="toctree-l2"><a class="reference internal" href="quickstart.html">Quickstart</a></li> <li class="toctree-l2"><a class="reference internal" href="key-concepts.html">Key concepts</a></li> <li class="toctree-l2"><a class="reference internal" href="jit-compilation.html">Just-in-time compilation</a></li> <li class="toctree-l2"><a class="reference internal" href="automatic-vectorization.html">Automatic vectorization</a></li> <li class="toctree-l2"><a class="reference internal" href="automatic-differentiation.html">Automatic differentiation</a></li> <li class="toctree-l2"><a class="reference internal" href="debugging.html">Introduction to debugging</a></li> <li class="toctree-l2"><a class="reference internal" href="random-numbers.html">Pseudorandom numbers</a></li> <li class="toctree-l2"><a class="reference internal" href="working-with-pytrees.html">Working with pytrees</a></li> <li class="toctree-l2"><a class="reference internal" href="sharded-computation.html">Introduction to parallel programming</a></li> <li class="toctree-l2"><a class="reference internal" href="stateful-computations.html">Stateful computations</a></li> <li class="toctree-l2"><a class="reference internal" href="control-flow.html">Control flow and logical operators with JIT</a></li> <li class="toctree-l2"><a class="reference internal" href="advanced-autodiff.html">Advanced automatic differentiation</a></li> <li class="toctree-l2"><a class="reference internal" href="external-callbacks.html">External callbacks</a></li> <li class="toctree-l2"><a class="reference internal" href="gradient-checkpointing.html">Gradient checkpointing with <code class="docutils literal notranslate">jax.checkpoint</code> (<code class="docutils literal notranslate">jax.remat</code>)</a></li> <li class="toctree-l2"><a class="reference internal" href="jax-primitives.html">JAX Internals: primitives</a></li> <li class="toctree-l2"><a class="reference internal" href="jaxpr.html">JAX internals: The jaxpr language</a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="notebooks/Common_Gotchas_in_JAX.html">🔪 JAX - The Sharp Bits 🔪</a></li> <li class="toctree-l1"><a class="reference internal" href="faq.html">Frequently asked questions (FAQ)</a></li> </ul> More guides/resources <ul class="current nav bd-sidenav"> <li class="toctree-l1 has-children"><a class="reference internal" href="user_guides.html">User guides</a><input class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-2"></label><ul> <li class="toctree-l2"><a class="reference internal" href="notebooks/thinking_in_jax.html">How to think in JAX</a></li> <li class="toctree-l2"><a class="reference internal" href="profiling.html">Profiling computation</a></li> <li class="toctree-l2"><a class="reference internal" href="device_memory_profiling.html">Profiling device memory</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="debugging/index.html">Debugging runtime values</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-3"></label><ul> <li class="toctree-l3"><a class="reference internal" href="debugging/print_breakpoint.html">Compiled prints and breakpoints</a></li> <li class="toctree-l3"><a class="reference internal" href="debugging/checkify_guide.html">The <code class="docutils literal notranslate">checkify</code> transformation</a></li> <li class="toctree-l3"><a class="reference internal" href="debugging/flags.html">JAX debugging flags</a></li> </ul> </li> <li class="toctree-l2"><a class="reference internal" href="gpu_performance_tips.html">GPU performance tips</a></li> <li class="toctree-l2"><a class="reference internal" href="persistent_compilation_cache.html">Persistent compilation cache</a></li> <li class="toctree-l2"><a class="reference internal" href="pytrees.html">Pytrees</a></li> <li class="toctree-l2"><a class="reference internal" href="errors.html">Errors</a></li> <li class="toctree-l2"><a class="reference internal" href="aot.html">Ahead-of-time lowering and compilation</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="export/index.html">Exporting and serialization</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-4"></label><ul> <li class="toctree-l3"><a class="reference internal" href="export/export.html">Exporting and serializing staged-out computations</a></li> <li class="toctree-l3"><a class="reference internal" href="export/shape_poly.html">Shape polymorphism</a></li> <li class="toctree-l3"><a class="reference internal" href="export/jax2tf.html">Interoperation with TensorFlow</a></li> </ul> </li> <li class="toctree-l2"><a class="reference internal" href="type_promotion.html">Type promotion semantics</a></li> <li class="toctree-l2"><a class="reference internal" href="transfer_guard.html">Transfer guard</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="pallas/index.html">Pallas: a JAX kernel language</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-5"></label><ul> <li class="toctree-l3"><a class="reference internal" href="pallas/quickstart.html">Pallas Quickstart</a></li> <li class="toctree-l3"><a class="reference internal" href="pallas/grid_blockspec.html">Grids and BlockSpecs</a></li> <li class="toctree-l3 has-children"><a class="reference internal" href="pallas/tpu/index.html">Pallas TPU</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-6"></label><ul> <li class="toctree-l4"><a class="reference internal" href="pallas/tpu/details.html">Writing TPU kernels with Pallas</a></li> <li class="toctree-l4"><a class="reference internal" href="pallas/tpu/pipelining.html">Pipelining</a></li> <li class="toctree-l4"><a class="reference internal" href="pallas/tpu/matmul.html">Matrix Multiplication</a></li> <li class="toctree-l4"><a class="reference internal" href="pallas/tpu/sparse.html">Scalar Prefetch and Block-Sparse Computation</a></li> <li class="toctree-l4"><a class="reference internal" href="pallas/tpu/distributed.html">Distributed Computing in Pallas for TPUs</a></li> </ul> </li> <li class="toctree-l3 has-children"><a class="reference internal" href="pallas/design/index.html">Pallas Design Notes</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-7"></label><ul> <li class="toctree-l4"><a class="reference internal" href="pallas/design/design.html">Pallas Design</a></li> <li class="toctree-l4"><a class="reference internal" href="pallas/design/async_note.html">Pallas Async Operations</a></li> </ul> </li> <li class="toctree-l3"><a class="reference internal" href="pallas/CHANGELOG.html">Pallas Changelog</a></li> </ul> </li> <li class="toctree-l2"><a class="reference internal" href="ffi.html">Foreign function interface (FFI)</a></li> <li class="toctree-l2"><a class="reference internal" href="notebooks/neural_network_with_tfds_data.html">Training a simple neural network, with tensorflow/datasets data loading</a></li> <li class="toctree-l2"><a class="reference internal" href="notebooks/Neural_Network_and_Data_Loading.html">Training a simple neural network, with PyTorch data loading</a></li> <li class="toctree-l2"><a class="reference internal" href="notebooks/vmapped_log_probs.html">Autobatching for Bayesian inference</a></li> </ul> </li> <li class="toctree-l1 has-children"><a class="reference internal" href="advanced_guide.html">Advanced guides</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-8"></label><ul> <li class="toctree-l2"><a class="reference internal" href="notebooks/Distributed_arrays_and_automatic_parallelization.html">Distributed arrays and automatic parallelization</a></li> <li class="toctree-l2"><a class="reference internal" href="notebooks/shard_map.html">Manual parallelism with <code class="docutils literal notranslate">shard_map</code></a></li> <li class="toctree-l2"><a class="reference internal" href="multi_process.html">Multi-host and multi-process environments</a></li> <li class="toctree-l2"><a class="reference internal" href="distributed_data_loading.html">Distributed data loading</a></li> <li class="toctree-l2"><a class="reference internal" href="notebooks/autodiff_cookbook.html">The Autodiff Cookbook</a></li> <li class="toctree-l2"><a class="reference internal" href="notebooks/Custom_derivative_rules_for_Python_code.html">Custom derivative rules</a></li> <li class="toctree-l2"><a class="reference internal" href="notebooks/autodiff_remat.html">Control autodiff’s saved values with <code class="docutils literal notranslate">jax.checkpoint</code> (aka <code class="docutils literal notranslate">jax.remat</code>)</a></li> <li class="toctree-l2"><a class="reference internal" href="notebooks/convolutions.html">Generalized convolutions in JAX</a></li> <li class="toctree-l2"><a class="reference internal" href="xla_flags.html">List of XLA compiler flags</a></li> </ul> </li> <li class="toctree-l1 has-children"><a class="reference internal" href="contributor_guide.html">Developer notes</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-9"></label><ul> <li class="toctree-l2"><a class="reference internal" href="contributing.html">Contributing to JAX</a></li> <li class="toctree-l2"><a class="reference internal" href="developer.html">Building from source</a></li> <li class="toctree-l2"><a class="reference internal" href="investigating_a_regression.html">Investigating a regression</a></li> <li class="toctree-l2"><a class="reference internal" href="autodidax.html">Autodidax: JAX core from scratch</a></li> <li class="toctree-l2"><a class="reference internal" href="autodidax2_part1.html">Autodidax2, part 1: JAX from scratch, again</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="jep/index.html">JAX Enhancement Proposals (JEPs)</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-10"></label><ul> <li class="toctree-l3"><a class="reference internal" href="jep/263-prng.html">263: JAX PRNG Design</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/2026-custom-derivatives.html">2026: Custom JVP/VJP rules for JAX-transformable functions</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/4008-custom-vjp-update.html">4008: Custom VJP and `nondiff_argnums` update</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/4410-omnistaging.html">4410: Omnistaging</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/9263-typed-keys.html">9263: Typed keys & pluggable RNGs</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/9407-type-promotion.html">9407: Design of Type Promotion Semantics for JAX</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/9419-jax-versioning.html">9419: Jax and Jaxlib versioning</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/10657-sequencing-effects.html">10657: Sequencing side-effects in JAX</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/11830-new-remat-checkpoint.html">11830: `jax.remat` / `jax.checkpoint` new implementation</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/12049-type-annotations.html">12049: Type Annotation Roadmap for JAX</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/14273-shard-map.html">14273: `shard_map` (`shmap`) for simple per-device code</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/15856-jex.html">15856: `jax.extend`, an extensions module</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/17111-shmap-transpose.html">17111: Efficient transposition of `shard_map` (and other maps)</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/18137-numpy-scipy-scope.html">18137: Scope of JAX NumPy & SciPy Wrappers</a></li> <li class="toctree-l3"><a class="reference internal" href="jep/25516-effver.html">25516: Effort-based versioning</a></li> </ul> </li> </ul> </li> <li class="toctree-l1 has-children"><a class="reference internal" href="extensions.html">Extension guides</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-11"></label><ul> <li class="toctree-l2"><a class="reference internal" href="notebooks/Writing_custom_interpreters_in_Jax.html">Writing custom Jaxpr interpreters in JAX</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="jax.extend.html"><code class="docutils literal notranslate">jax.extend</code> module</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-12"></label><ul> <li class="toctree-l3"><a class="reference internal" href="jax.extend.core.html"><code class="docutils literal notranslate">jax.extend.core</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.extend.linear_util.html"><code class="docutils literal notranslate">jax.extend.linear_util</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.extend.mlir.html"><code class="docutils literal notranslate">jax.extend.mlir</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.extend.random.html"><code class="docutils literal notranslate">jax.extend.random</code> module</a></li> </ul> </li> <li class="toctree-l2"><a class="reference internal" href="building_on_jax.html">Building on JAX</a></li> </ul> </li> <li class="toctree-l1 has-children"><a class="reference internal" href="notes.html">Notes</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-13"></label><ul> <li class="toctree-l2"><a class="reference internal" href="api_compatibility.html">API compatibility</a></li> <li class="toctree-l2"><a class="reference internal" href="deprecation.html">Python and NumPy version support policy</a></li> <li class="toctree-l2"><a class="reference internal" href="jax_array_migration.html">jax.Array migration</a></li> <li class="toctree-l2"><a class="reference internal" href="async_dispatch.html">Asynchronous dispatch</a></li> <li class="toctree-l2"><a class="reference internal" href="concurrency.html">Concurrency</a></li> <li class="toctree-l2"><a class="reference internal" href="gpu_memory_allocation.html">GPU memory allocation</a></li> <li class="toctree-l2"><a class="reference internal" href="rank_promotion_warning.html">Rank promotion warning</a></li> </ul> </li> <li class="toctree-l1 current active has-children"><a class="reference internal" href="jax.html">Public API: <code class="docutils literal notranslate">jax</code> package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-14"></label><ul class="current"> <li class="toctree-l2 has-children"><a class="reference internal" href="jax.numpy.html"><code class="docutils literal notranslate">jax.numpy</code> module</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-15"></label><ul> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.fft.html">jax.numpy.fft.fft</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.fft2.html">jax.numpy.fft.fft2</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.fftfreq.html">jax.numpy.fft.fftfreq</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.fftn.html">jax.numpy.fft.fftn</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.fftshift.html">jax.numpy.fft.fftshift</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.hfft.html">jax.numpy.fft.hfft</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.ifft.html">jax.numpy.fft.ifft</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.ifft2.html">jax.numpy.fft.ifft2</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.ifftn.html">jax.numpy.fft.ifftn</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.ifftshift.html">jax.numpy.fft.ifftshift</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.ihfft.html">jax.numpy.fft.ihfft</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.irfft.html">jax.numpy.fft.irfft</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.irfft2.html">jax.numpy.fft.irfft2</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.irfftn.html">jax.numpy.fft.irfftn</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.rfft.html">jax.numpy.fft.rfft</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.rfft2.html">jax.numpy.fft.rfft2</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.rfftfreq.html">jax.numpy.fft.rfftfreq</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.numpy.fft.rfftn.html">jax.numpy.fft.rfftn</a></li> </ul> </li> <li class="toctree-l2 has-children"><a class="reference internal" href="jax.scipy.html"><code class="docutils literal notranslate">jax.scipy</code> module</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-16"></label><ul> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.scipy.stats.bernoulli.logpmf.html">jax.scipy.stats.bernoulli.logpmf</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.scipy.stats.bernoulli.pmf.html">jax.scipy.stats.bernoulli.pmf</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.scipy.stats.bernoulli.cdf.html">jax.scipy.stats.bernoulli.cdf</a></li> <li class="toctree-l3"><a class="reference internal" href="_autosummary/jax.scipy.stats.bernoulli.ppf.html">jax.scipy.stats.bernoulli.ppf</a></li> </ul> </li> <li class="toctree-l2"><a class="reference internal" href="jax.lax.html"><code class="docutils literal notranslate">jax.lax</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.random.html"><code class="docutils literal notranslate">jax.random</code> module</a></li> <li class="toctree-l2 current active"><a class="current reference internal" href="#"><code class="docutils literal notranslate">jax.sharding</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.debug.html"><code class="docutils literal notranslate">jax.debug</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.dlpack.html"><code class="docutils literal notranslate">jax.dlpack</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.distributed.html"><code class="docutils literal notranslate">jax.distributed</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.dtypes.html"><code class="docutils literal notranslate">jax.dtypes</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.ffi.html"><code class="docutils literal notranslate">jax.ffi</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.flatten_util.html"><code class="docutils literal notranslate">jax.flatten_util</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.image.html"><code class="docutils literal notranslate">jax.image</code> module</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="jax.nn.html"><code class="docutils literal notranslate">jax.nn</code> module</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-17"></label><ul> <li class="toctree-l3"><a class="reference internal" href="jax.nn.initializers.html"><code class="docutils literal notranslate">jax.nn.initializers</code> module</a></li> </ul> </li> <li class="toctree-l2"><a class="reference internal" href="jax.ops.html"><code class="docutils literal notranslate">jax.ops</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.profiler.html"><code class="docutils literal notranslate">jax.profiler</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.stages.html"><code class="docutils literal notranslate">jax.stages</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.test_util.html"><code class="docutils literal notranslate">jax.test_util</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.tree.html"><code class="docutils literal notranslate">jax.tree</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.tree_util.html"><code class="docutils literal notranslate">jax.tree_util</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.typing.html"><code class="docutils literal notranslate">jax.typing</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="jax.export.html"><code class="docutils literal notranslate">jax.export</code> module</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="jax.extend.html"><code class="docutils literal notranslate">jax.extend</code> module</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-18"></label><ul> <li class="toctree-l3"><a class="reference internal" href="jax.extend.core.html"><code class="docutils literal notranslate">jax.extend.core</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.extend.linear_util.html"><code class="docutils literal notranslate">jax.extend.linear_util</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.extend.mlir.html"><code class="docutils literal notranslate">jax.extend.mlir</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.extend.random.html"><code class="docutils literal notranslate">jax.extend.random</code> module</a></li> </ul> </li> <li class="toctree-l2 has-children"><a class="reference internal" href="jax.example_libraries.html"><code class="docutils literal notranslate">jax.example_libraries</code> module</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-19"></label><ul> <li class="toctree-l3"><a class="reference internal" href="jax.example_libraries.optimizers.html"><code class="docutils literal notranslate">jax.example_libraries.optimizers</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.example_libraries.stax.html"><code class="docutils literal notranslate">jax.example_libraries.stax</code> module</a></li> </ul> </li> <li class="toctree-l2 has-children"><a class="reference internal" href="jax.experimental.html"><code class="docutils literal notranslate">jax.experimental</code> module</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-20"></label><ul> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.checkify.html"><code class="docutils literal notranslate">jax.experimental.checkify</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.compilation_cache.html"><code class="docutils literal notranslate">jax.experimental.compilation_cache</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.custom_dce.html"><code class="docutils literal notranslate">jax.experimental.custom_dce</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.custom_partitioning.html"><code class="docutils literal notranslate">jax.experimental.custom_partitioning</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.jet.html"><code class="docutils literal notranslate">jax.experimental.jet</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.key_reuse.html"><code class="docutils literal notranslate">jax.experimental.key_reuse</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.mesh_utils.html"><code class="docutils literal notranslate">jax.experimental.mesh_utils</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.multihost_utils.html"><code class="docutils literal notranslate">jax.experimental.multihost_utils</code> module</a></li> <li class="toctree-l3 has-children"><a class="reference internal" href="jax.experimental.pallas.html"><code class="docutils literal notranslate">jax.experimental.pallas</code> module</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-21"></label><ul> <li class="toctree-l4"><a class="reference internal" href="jax.experimental.pallas.mosaic_gpu.html"><code class="docutils literal notranslate">jax.experimental.pallas.mosaic_gpu</code> module</a></li> <li class="toctree-l4"><a class="reference internal" href="jax.experimental.pallas.triton.html"><code class="docutils literal notranslate">jax.experimental.pallas.triton</code> module</a></li> <li class="toctree-l4"><a class="reference internal" href="jax.experimental.pallas.tpu.html"><code class="docutils literal notranslate">jax.experimental.pallas.tpu</code> module</a></li> </ul> </li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.pjit.html"><code class="docutils literal notranslate">jax.experimental.pjit</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.serialize_executable.html"><code class="docutils literal notranslate">jax.experimental.serialize_executable</code> module</a></li> <li class="toctree-l3"><a class="reference internal" href="jax.experimental.shard_map.html"><code class="docutils literal notranslate">jax.experimental.shard_map</code> module</a></li> <li class="toctree-l3 has-children"><a class="reference internal" href="jax.experimental.sparse.html"><code class="docutils literal notranslate">jax.experimental.sparse</code> module</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-22"></label><ul> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.BCOO.html">jax.experimental.sparse.BCOO</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_broadcast_in_dim.html">jax.experimental.sparse.bcoo_broadcast_in_dim</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_concatenate.html">jax.experimental.sparse.bcoo_concatenate</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_dot_general.html">jax.experimental.sparse.bcoo_dot_general</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_dot_general_sampled.html">jax.experimental.sparse.bcoo_dot_general_sampled</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_dynamic_slice.html">jax.experimental.sparse.bcoo_dynamic_slice</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_extract.html">jax.experimental.sparse.bcoo_extract</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_fromdense.html">jax.experimental.sparse.bcoo_fromdense</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_gather.html">jax.experimental.sparse.bcoo_gather</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_multiply_dense.html">jax.experimental.sparse.bcoo_multiply_dense</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_multiply_sparse.html">jax.experimental.sparse.bcoo_multiply_sparse</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_update_layout.html">jax.experimental.sparse.bcoo_update_layout</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_reduce_sum.html">jax.experimental.sparse.bcoo_reduce_sum</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_reshape.html">jax.experimental.sparse.bcoo_reshape</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_slice.html">jax.experimental.sparse.bcoo_slice</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_sort_indices.html">jax.experimental.sparse.bcoo_sort_indices</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_squeeze.html">jax.experimental.sparse.bcoo_squeeze</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_sum_duplicates.html">jax.experimental.sparse.bcoo_sum_duplicates</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_todense.html">jax.experimental.sparse.bcoo_todense</a></li> <li class="toctree-l4"><a class="reference internal" href="_autosummary/jax.experimental.sparse.bcoo_transpose.html">jax.experimental.sparse.bcoo_transpose</a></li> </ul> </li> </ul> </li> <li class="toctree-l2"><a class="reference internal" href="jax.lib.html"><code class="docutils literal notranslate">jax.lib</code> module</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.addressable_shards.html">jax.Array.addressable_shards</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.all.html">jax.Array.all</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.any.html">jax.Array.any</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.argmax.html">jax.Array.argmax</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.argmin.html">jax.Array.argmin</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.argpartition.html">jax.Array.argpartition</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.argsort.html">jax.Array.argsort</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.astype.html">jax.Array.astype</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.at.html">jax.Array.at</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.choose.html">jax.Array.choose</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.clip.html">jax.Array.clip</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.compress.html">jax.Array.compress</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.committed.html">jax.Array.committed</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.conj.html">jax.Array.conj</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.conjugate.html">jax.Array.conjugate</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.copy.html">jax.Array.copy</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.copy_to_host_async.html">jax.Array.copy_to_host_async</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.cumprod.html">jax.Array.cumprod</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.cumsum.html">jax.Array.cumsum</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.device.html">jax.Array.device</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.diagonal.html">jax.Array.diagonal</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.dot.html">jax.Array.dot</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.dtype.html">jax.Array.dtype</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.flat.html">jax.Array.flat</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.flatten.html">jax.Array.flatten</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.global_shards.html">jax.Array.global_shards</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.imag.html">jax.Array.imag</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.is_fully_addressable.html">jax.Array.is_fully_addressable</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.is_fully_replicated.html">jax.Array.is_fully_replicated</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.item.html">jax.Array.item</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.itemsize.html">jax.Array.itemsize</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.max.html">jax.Array.max</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.mean.html">jax.Array.mean</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.min.html">jax.Array.min</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.nbytes.html">jax.Array.nbytes</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.ndim.html">jax.Array.ndim</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.nonzero.html">jax.Array.nonzero</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.prod.html">jax.Array.prod</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.ptp.html">jax.Array.ptp</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.ravel.html">jax.Array.ravel</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.real.html">jax.Array.real</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.repeat.html">jax.Array.repeat</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.reshape.html">jax.Array.reshape</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.round.html">jax.Array.round</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.searchsorted.html">jax.Array.searchsorted</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.shape.html">jax.Array.shape</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.sharding.html">jax.Array.sharding</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.size.html">jax.Array.size</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.sort.html">jax.Array.sort</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.squeeze.html">jax.Array.squeeze</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.std.html">jax.Array.std</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.sum.html">jax.Array.sum</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.swapaxes.html">jax.Array.swapaxes</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.take.html">jax.Array.take</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.to_device.html">jax.Array.to_device</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.trace.html">jax.Array.trace</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.transpose.html">jax.Array.transpose</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.var.html">jax.Array.var</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.view.html">jax.Array.view</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.T.html">jax.Array.T</a></li> <li class="toctree-l2"><a class="reference internal" href="_autosummary/jax.Array.mT.html">jax.Array.mT</a></li> </ul> </li> <li class="toctree-l1"><a class="reference internal" href="about.html">About the project</a></li> </ul> <ul class="nav bd-sidenav"> <li class="toctree-l1"><a class="reference internal" href="changelog.html">Change log</a></li> <li class="toctree-l1"><a class="reference internal" href="glossary.html">Glossary of terms</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"> <div class="sbt-scroll-pixel-helper"></div> <div class="bd-content"> <div class="bd-article-container"> <div class="bd-header-article"> <div class="header-article-items header-article__inner"> <div class="header-article-items__start"> <div class="header-article-item"><label class="sidebar-toggle primary-toggle btn btn-sm" for="__primary" title="Toggle primary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip"> </label></div> </div> <div class="header-article-items__end"> <div class="header-article-item"> <div class="article-header-buttons"> <a href="https://github.com/jax-ml/jax" 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/jax.sharding.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 navbar-btn 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 navbar-btn search-button search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip"> </button> `); </script> <label class="sidebar-toggle secondary-toggle btn btn-sm" for="__secondary"title="Toggle secondary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip"> </label> </div></div> </div> </div> </div> <div id="jb-print-docs-body" class="onlyprint"> <h1>jax.sharding module</h1>  <div id="print-main-content"> <div id="jb-print-toc"> <div> <h2> Contents </h2> </div> <nav aria-label="Page"> <ul class="visible nav section-nav flex-column"> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#classes">Classes</a><ul class="visible nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding"><code class="docutils literal notranslate">Sharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.addressable_devices"><code class="docutils literal notranslate">Sharding.addressable_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.addressable_devices_indices_map"><code class="docutils literal notranslate">Sharding.addressable_devices_indices_map()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.device_set"><code class="docutils literal notranslate">Sharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.devices_indices_map"><code class="docutils literal notranslate">Sharding.devices_indices_map()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.is_equivalent_to"><code class="docutils literal notranslate">Sharding.is_equivalent_to()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.is_fully_addressable"><code class="docutils literal notranslate">Sharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.is_fully_replicated"><code class="docutils literal notranslate">Sharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.memory_kind"><code class="docutils literal notranslate">Sharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.num_devices"><code class="docutils literal notranslate">Sharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.shard_shape"><code class="docutils literal notranslate">Sharding.shard_shape()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.with_memory_kind"><code class="docutils literal notranslate">Sharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding"><code class="docutils literal notranslate">SingleDeviceSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.device_set"><code class="docutils literal notranslate">SingleDeviceSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.devices_indices_map"><code class="docutils literal notranslate">SingleDeviceSharding.devices_indices_map()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.is_fully_addressable"><code class="docutils literal notranslate">SingleDeviceSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.is_fully_replicated"><code class="docutils literal notranslate">SingleDeviceSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.memory_kind"><code class="docutils literal notranslate">SingleDeviceSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.num_devices"><code class="docutils literal notranslate">SingleDeviceSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.with_memory_kind"><code class="docutils literal notranslate">SingleDeviceSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding"><code class="docutils literal notranslate">NamedSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.addressable_devices"><code class="docutils literal notranslate">NamedSharding.addressable_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.device_set"><code class="docutils literal notranslate">NamedSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.is_fully_addressable"><code class="docutils literal notranslate">NamedSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.is_fully_replicated"><code class="docutils literal notranslate">NamedSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.memory_kind"><code class="docutils literal notranslate">NamedSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.mesh"><code class="docutils literal notranslate">NamedSharding.mesh</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.num_devices"><code class="docutils literal notranslate">NamedSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.spec"><code class="docutils literal notranslate">NamedSharding.spec</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.with_memory_kind"><code class="docutils literal notranslate">NamedSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding"><code class="docutils literal notranslate">PositionalSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.device_set"><code class="docutils literal notranslate">PositionalSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.is_fully_addressable"><code class="docutils literal notranslate">PositionalSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.is_fully_replicated"><code class="docutils literal notranslate">PositionalSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.memory_kind"><code class="docutils literal notranslate">PositionalSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.num_devices"><code class="docutils literal notranslate">PositionalSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.with_memory_kind"><code class="docutils literal notranslate">PositionalSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding"><code class="docutils literal notranslate">PmapSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.default"><code class="docutils literal notranslate">PmapSharding.default()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.device_set"><code class="docutils literal notranslate">PmapSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.devices"><code class="docutils literal notranslate">PmapSharding.devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.devices_indices_map"><code class="docutils literal notranslate">PmapSharding.devices_indices_map()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.is_equivalent_to"><code class="docutils literal notranslate">PmapSharding.is_equivalent_to()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.is_fully_addressable"><code class="docutils literal notranslate">PmapSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.is_fully_replicated"><code class="docutils literal notranslate">PmapSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.memory_kind"><code class="docutils literal notranslate">PmapSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.num_devices"><code class="docutils literal notranslate">PmapSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.shard_shape"><code class="docutils literal notranslate">PmapSharding.shard_shape()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.sharding_spec"><code class="docutils literal notranslate">PmapSharding.sharding_spec</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.with_memory_kind"><code class="docutils literal notranslate">PmapSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding"><code class="docutils literal notranslate">GSPMDSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.device_set"><code class="docutils literal notranslate">GSPMDSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.is_fully_addressable"><code class="docutils literal notranslate">GSPMDSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.is_fully_replicated"><code class="docutils literal notranslate">GSPMDSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.memory_kind"><code class="docutils literal notranslate">GSPMDSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.num_devices"><code class="docutils literal notranslate">GSPMDSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.with_memory_kind"><code class="docutils literal notranslate">GSPMDSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PartitionSpec"><code class="docutils literal notranslate">PartitionSpec</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Mesh"><code class="docutils literal notranslate">Mesh</code></a></li> </ul> </li> </ul> </nav> </div> </div> </div> <div id="searchbox"></div> <article class="bd-article" role="main"> <section id="module-jax.sharding"> <h1><code class="docutils literal notranslate">jax.sharding</code> module<a class="headerlink" href="#module-jax.sharding" title="Link to this heading">#</a></h1> <section id="classes"> <h2>Classes<a class="headerlink" href="#classes" title="Link to this heading">#</a></h2> <dl class="py class"> <dt class="sig sig-object py" id="jax.sharding.Sharding"> class jax.sharding.Sharding<a class="headerlink" href="#jax.sharding.Sharding" title="Link to this definition">#</a></dt> <dd>Describes how a <a class="reference internal" href="_autosummary/jax.Array.html#jax.Array" title="jax.Array"><code class="xref py py-class docutils literal notranslate">jax.Array</code></a> is laid out across devices. <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.Sharding.addressable_devices"> property addressable_devices: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#set" title="(in Python v3.13)">set</a>[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>]<a class="headerlink" href="#jax.sharding.Sharding.addressable_devices" title="Link to this definition">#</a></dt> <dd>The set of devices in the <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> that are addressable by the current process. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.Sharding.addressable_devices_indices_map"> addressable_devices_indices_map(global_shape)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L156-L164">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.addressable_devices_indices_map" title="Link to this definition">#</a></dt> <dd>A mapping from addressable devices to the slice of array data each contains. <code class="docutils literal notranslate">addressable_devices_indices_map</code> contains that part of <code class="docutils literal notranslate">device_indices_map</code> that applies to the addressable devices. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">global_shape (Shape) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even">Mapping[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>, Index | None] </dd> </dl> </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.Sharding.device_set"> property device_set: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#set" title="(in Python v3.13)">set</a>[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>]<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L86-L94">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.device_set" title="Link to this definition">#</a></dt> <dd>The set of devices that this <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> spans. In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.Sharding.devices_indices_map"> devices_indices_map(global_shape)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L165-L172">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.devices_indices_map" title="Link to this definition">#</a></dt> <dd>Returns a mapping from devices to the array slices each contains. The mapping includes all global devices, i.e., including non-addressable devices from other processes. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">global_shape (Shape) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even">Mapping[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>, Index] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.Sharding.is_equivalent_to"> is_equivalent_to(other, ndim)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L190-L210">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.is_equivalent_to" title="Link to this definition">#</a></dt> <dd>Returns <code class="docutils literal notranslate">True</code> if two shardings are equivalent. Two shardings are equivalent if they place the same logical array shards on the same devices. For example, a <a class="reference internal" href="#jax.sharding.NamedSharding" title="jax.sharding.NamedSharding"><code class="xref py py-class docutils literal notranslate">NamedSharding</code></a> may be equivalent to a <a class="reference internal" href="#jax.sharding.PositionalSharding" title="jax.sharding.PositionalSharding"><code class="xref py py-class docutils literal notranslate">PositionalSharding</code></a> if both place the same shards of the array on the same devices. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd"><ul class="simple"> <li>self (<a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding">Sharding</a>)</li> <li>other (<a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding">Sharding</a>)</li> <li>ndim (<a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)">int</a>)</li> </ul> </dd> <dt class="field-even">Return type:</dt> <dd class="field-even"><a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a> </dd> </dl> </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.Sharding.is_fully_addressable"> property is_fully_addressable: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L104-L113">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.is_fully_addressable" title="Link to this definition">#</a></dt> <dd>Is this sharding fully addressable? A sharding is fully addressable if the current process can address all of the devices named in the <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a>. <code class="docutils literal notranslate">is_fully_addressable</code> is equivalent to “is_local” in multi-process JAX. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.Sharding.is_fully_replicated"> property is_fully_replicated: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L95-L103">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.is_fully_replicated" title="Link to this definition">#</a></dt> <dd>Is this sharding fully replicated? A sharding is fully replicated if each device has a complete copy of the entire data. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.Sharding.memory_kind"> property memory_kind: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a> | <a class="reference external" href="https://docs.python.org/3/library/constants.html#None" title="(in Python v3.13)">None</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L119-L123">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.memory_kind" title="Link to this definition">#</a></dt> <dd>Returns the memory kind of the sharding. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.Sharding.num_devices"> property num_devices: <a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)">int</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L114-L118">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.num_devices" title="Link to this definition">#</a></dt> <dd>Number of devices that the sharding contains. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.Sharding.shard_shape"> shard_shape(global_shape)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L182-L189">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.shard_shape" title="Link to this definition">#</a></dt> <dd>Returns the shape of the data on each device. The shard shape returned by this function is calculated from <code class="docutils literal notranslate">global_shape</code> and the properties of the sharding. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">global_shape (Shape) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even">Shape </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.Sharding.with_memory_kind"> with_memory_kind(kind)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding.py#L124-L127">[source]</a><a class="headerlink" href="#jax.sharding.Sharding.with_memory_kind" title="Link to this definition">#</a></dt> <dd>Returns a new Sharding instance with the specified memory kind. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">kind (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a>) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even"><a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding">Sharding</a> </dd> </dl> </dd></dl> </dd></dl> <dl class="py class"> <dt class="sig sig-object py" id="jax.sharding.SingleDeviceSharding"> class jax.sharding.SingleDeviceSharding<a class="headerlink" href="#jax.sharding.SingleDeviceSharding" title="Link to this definition">#</a></dt> <dd>Bases: <a class="reference internal" href="#jax.sharding.Sharding" title="jaxlib.xla_extension.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> A <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> that places its data on a single device. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">device – A single <code class="xref py py-class docutils literal notranslate">Device</code>. </dd> </dl> Examples <div class="doctest highlight-default notranslate"><div class="highlight"><pre>>>> single_device_sharding = jax.sharding.SingleDeviceSharding( ... jax.devices()[0]) </pre></div> </div> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.SingleDeviceSharding.device_set"> property device_set: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#set" title="(in Python v3.13)">set</a>[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>]<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L165-L168">[source]</a><a class="headerlink" href="#jax.sharding.SingleDeviceSharding.device_set" title="Link to this definition">#</a></dt> <dd>The set of devices that this <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> spans. In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.SingleDeviceSharding.devices_indices_map"> devices_indices_map(global_shape)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L176-L178">[source]</a><a class="headerlink" href="#jax.sharding.SingleDeviceSharding.devices_indices_map" title="Link to this definition">#</a></dt> <dd>Returns a mapping from devices to the array slices each contains. The mapping includes all global devices, i.e., including non-addressable devices from other processes. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">global_shape (Shape) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even">Mapping[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>, Index] </dd> </dl> </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.SingleDeviceSharding.is_fully_addressable"> property is_fully_addressable: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L195-L198">[source]</a><a class="headerlink" href="#jax.sharding.SingleDeviceSharding.is_fully_addressable" title="Link to this definition">#</a></dt> <dd>Is this sharding fully addressable? A sharding is fully addressable if the current process can address all of the devices named in the <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a>. <code class="docutils literal notranslate">is_fully_addressable</code> is equivalent to “is_local” in multi-process JAX. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.SingleDeviceSharding.is_fully_replicated"> property is_fully_replicated: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L191-L194">[source]</a><a class="headerlink" href="#jax.sharding.SingleDeviceSharding.is_fully_replicated" title="Link to this definition">#</a></dt> <dd>Is this sharding fully replicated? A sharding is fully replicated if each device has a complete copy of the entire data. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.SingleDeviceSharding.memory_kind"> property memory_kind: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a> | <a class="reference external" href="https://docs.python.org/3/library/constants.html#None" title="(in Python v3.13)">None</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L169-L172">[source]</a><a class="headerlink" href="#jax.sharding.SingleDeviceSharding.memory_kind" title="Link to this definition">#</a></dt> <dd>Returns the memory kind of the sharding. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.SingleDeviceSharding.num_devices"> property num_devices: <a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)">int</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L161-L164">[source]</a><a class="headerlink" href="#jax.sharding.SingleDeviceSharding.num_devices" title="Link to this definition">#</a></dt> <dd>Number of devices that the sharding contains. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.SingleDeviceSharding.with_memory_kind"> with_memory_kind(kind)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L173-L175">[source]</a><a class="headerlink" href="#jax.sharding.SingleDeviceSharding.with_memory_kind" title="Link to this definition">#</a></dt> <dd>Returns a new Sharding instance with the specified memory kind. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">kind (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a>) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even"><a class="reference internal" href="#jax.sharding.SingleDeviceSharding" title="jax.sharding.SingleDeviceSharding">SingleDeviceSharding</a> </dd> </dl> </dd></dl> </dd></dl> <dl class="py class"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding"> class jax.sharding.NamedSharding<a class="headerlink" href="#jax.sharding.NamedSharding" title="Link to this definition">#</a></dt> <dd>Bases: <a class="reference internal" href="#jax.sharding.Sharding" title="jaxlib.xla_extension.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> A <a class="reference internal" href="#jax.sharding.NamedSharding" title="jax.sharding.NamedSharding"><code class="xref py py-class docutils literal notranslate">NamedSharding</code></a> expresses sharding using named axes. A <a class="reference internal" href="#jax.sharding.NamedSharding" title="jax.sharding.NamedSharding"><code class="xref py py-class docutils literal notranslate">NamedSharding</code></a> is a pair of a <a class="reference internal" href="#jax.sharding.Mesh" title="jax.sharding.Mesh"><code class="xref py py-class docutils literal notranslate">Mesh</code></a> of devices and <a class="reference internal" href="#jax.sharding.PartitionSpec" title="jax.sharding.PartitionSpec"><code class="xref py py-class docutils literal notranslate">PartitionSpec</code></a> which describes how to shard an array across that mesh. A <a class="reference internal" href="#jax.sharding.Mesh" title="jax.sharding.Mesh"><code class="xref py py-class docutils literal notranslate">Mesh</code></a> is a multidimensional NumPy array of JAX devices, where each axis of the mesh has a name, e.g. <code class="docutils literal notranslate">'x'</code> or <code class="docutils literal notranslate">'y'</code>. A <a class="reference internal" href="#jax.sharding.PartitionSpec" title="jax.sharding.PartitionSpec"><code class="xref py py-class docutils literal notranslate">PartitionSpec</code></a> is a tuple, whose elements can be a <code class="docutils literal notranslate">None</code>, a mesh axis, or a tuple of mesh axes. Each element describes how an input dimension is partitioned across zero or more mesh dimensions. For example, <code class="docutils literal notranslate">PartitionSpec('x', 'y')</code> says that the first dimension of data is sharded across <code class="docutils literal notranslate">x</code> axis of the mesh, and the second dimension is sharded across <code class="docutils literal notranslate">y</code> axis of the mesh. The Distributed arrays and automatic parallelization (<a class="reference external" href="https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#namedsharding-gives-a-way-to-express-shardings-with-names">https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#namedsharding-gives-a-way-to-express-shardings-with-names</a>) tutorial has more details and diagrams that explain how <a class="reference internal" href="#jax.sharding.Mesh" title="jax.sharding.Mesh"><code class="xref py py-class docutils literal notranslate">Mesh</code></a> and <a class="reference internal" href="#jax.sharding.PartitionSpec" title="jax.sharding.PartitionSpec"><code class="xref py py-class docutils literal notranslate">PartitionSpec</code></a> are used. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd"><ul class="simple"> <li>mesh – A <a class="reference internal" href="#jax.sharding.Mesh" title="jax.sharding.Mesh"><code class="xref py py-class docutils literal notranslate">jax.sharding.Mesh</code></a> object.</li> <li>spec – A <a class="reference internal" href="#jax.sharding.PartitionSpec" title="jax.sharding.PartitionSpec"><code class="xref py py-class docutils literal notranslate">jax.sharding.PartitionSpec</code></a> object.</li> </ul> </dd> </dl> Examples <div class="doctest highlight-default notranslate"><div class="highlight"><pre>>>> from jax.sharding import Mesh >>> from jax.sharding import PartitionSpec as P >>> mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y')) >>> spec = P('x', 'y') >>> named_sharding = jax.sharding.NamedSharding(mesh, spec) </pre></div> </div> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding.addressable_devices"> property addressable_devices: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#set" title="(in Python v3.13)">set</a>[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>]<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/named_sharding.py#L222-L230">[source]</a><a class="headerlink" href="#jax.sharding.NamedSharding.addressable_devices" title="Link to this definition">#</a></dt> <dd>The set of devices in the <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> that are addressable by the current process. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding.device_set"> property device_set: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#set" title="(in Python v3.13)">set</a>[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>]<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/named_sharding.py#L193-L199">[source]</a><a class="headerlink" href="#jax.sharding.NamedSharding.device_set" title="Link to this definition">#</a></dt> <dd>The set of devices that this <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> spans. In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding.is_fully_addressable"> property is_fully_addressable: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/named_sharding.py#L207-L215">[source]</a><a class="headerlink" href="#jax.sharding.NamedSharding.is_fully_addressable" title="Link to this definition">#</a></dt> <dd>Is this sharding fully addressable? A sharding is fully addressable if the current process can address all of the devices named in the <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a>. <code class="docutils literal notranslate">is_fully_addressable</code> is equivalent to “is_local” in multi-process JAX. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding.is_fully_replicated"> property is_fully_replicated: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="headerlink" href="#jax.sharding.NamedSharding.is_fully_replicated" title="Link to this definition">#</a></dt> <dd>Is this sharding fully replicated? A sharding is fully replicated if each device has a complete copy of the entire data. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding.memory_kind"> property memory_kind: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a> | <a class="reference external" href="https://docs.python.org/3/library/constants.html#None" title="(in Python v3.13)">None</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/named_sharding.py#L146-L149">[source]</a><a class="headerlink" href="#jax.sharding.NamedSharding.memory_kind" title="Link to this definition">#</a></dt> <dd>Returns the memory kind of the sharding. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding.mesh"> property mesh<a class="headerlink" href="#jax.sharding.NamedSharding.mesh" title="Link to this definition">#</a></dt> <dd>(self) -> object </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding.num_devices"> property num_devices: <a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)">int</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/named_sharding.py#L189-L192">[source]</a><a class="headerlink" href="#jax.sharding.NamedSharding.num_devices" title="Link to this definition">#</a></dt> <dd>Number of devices that the sharding contains. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding.spec"> property spec<a class="headerlink" href="#jax.sharding.NamedSharding.spec" title="Link to this definition">#</a></dt> <dd>(self) -> object </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.NamedSharding.with_memory_kind"> with_memory_kind(kind)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/named_sharding.py#L242-L244">[source]</a><a class="headerlink" href="#jax.sharding.NamedSharding.with_memory_kind" title="Link to this definition">#</a></dt> <dd>Returns a new Sharding instance with the specified memory kind. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">kind (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a>) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even"><a class="reference internal" href="#jax.sharding.NamedSharding" title="jax.sharding.NamedSharding">NamedSharding</a> </dd> </dl> </dd></dl> </dd></dl> <dl class="py class"> <dt class="sig sig-object py" id="jax.sharding.PositionalSharding"> class jax.sharding.PositionalSharding(devices, *, memory_kind=None)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L419-L545">[source]</a><a class="headerlink" href="#jax.sharding.PositionalSharding" title="Link to this definition">#</a></dt> <dd>Bases: <a class="reference internal" href="#jax.sharding.Sharding" title="jaxlib.xla_extension.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd"><ul class="simple"> <li>devices (Sequence[xc.Device] | np.ndarray)</li> <li>memory_kind (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a> | None)</li> </ul> </dd> </dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PositionalSharding.device_set"> property device_set: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#set" title="(in Python v3.13)">set</a>[xc.Device]<a class="headerlink" href="#jax.sharding.PositionalSharding.device_set" title="Link to this definition">#</a></dt> <dd>The set of devices that this <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> spans. In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PositionalSharding.is_fully_addressable"> property is_fully_addressable: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="headerlink" href="#jax.sharding.PositionalSharding.is_fully_addressable" title="Link to this definition">#</a></dt> <dd>Is this sharding fully addressable? A sharding is fully addressable if the current process can address all of the devices named in the <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a>. <code class="docutils literal notranslate">is_fully_addressable</code> is equivalent to “is_local” in multi-process JAX. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PositionalSharding.is_fully_replicated"> property is_fully_replicated: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="headerlink" href="#jax.sharding.PositionalSharding.is_fully_replicated" title="Link to this definition">#</a></dt> <dd>Is this sharding fully replicated? A sharding is fully replicated if each device has a complete copy of the entire data. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PositionalSharding.memory_kind"> property memory_kind: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a> | <a class="reference external" href="https://docs.python.org/3/library/constants.html#None" title="(in Python v3.13)">None</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L518-L521">[source]</a><a class="headerlink" href="#jax.sharding.PositionalSharding.memory_kind" title="Link to this definition">#</a></dt> <dd>Returns the memory kind of the sharding. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PositionalSharding.num_devices"> property num_devices: <a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)">int</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L510-L513">[source]</a><a class="headerlink" href="#jax.sharding.PositionalSharding.num_devices" title="Link to this definition">#</a></dt> <dd>Number of devices that the sharding contains. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.PositionalSharding.with_memory_kind"> with_memory_kind(kind)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L522-L524">[source]</a><a class="headerlink" href="#jax.sharding.PositionalSharding.with_memory_kind" title="Link to this definition">#</a></dt> <dd>Returns a new Sharding instance with the specified memory kind. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">kind (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a>) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even"><a class="reference internal" href="#jax.sharding.PositionalSharding" title="jax.sharding.PositionalSharding">PositionalSharding</a> </dd> </dl> </dd></dl> </dd></dl> <dl class="py class"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding"> class jax.sharding.PmapSharding<a class="headerlink" href="#jax.sharding.PmapSharding" title="Link to this definition">#</a></dt> <dd>Bases: <a class="reference internal" href="#jax.sharding.Sharding" title="jaxlib.xla_extension.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> Describes a sharding used by <a class="reference internal" href="_autosummary/jax.pmap.html#jax.pmap" title="jax.pmap"><code class="xref py py-func docutils literal notranslate">jax.pmap()</code></a>. <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.default"> classmethod default(shape, sharded_dim=0, devices=None)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L256-L300">[source]</a><a class="headerlink" href="#jax.sharding.PmapSharding.default" title="Link to this definition">#</a></dt> <dd>Creates a <a class="reference internal" href="#jax.sharding.PmapSharding" title="jax.sharding.PmapSharding"><code class="xref py py-class docutils literal notranslate">PmapSharding</code></a> which matches the default placement used by <a class="reference internal" href="_autosummary/jax.pmap.html#jax.pmap" title="jax.pmap"><code class="xref py py-func docutils literal notranslate">jax.pmap()</code></a>. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd"><ul class="simple"> <li>shape (Shape) – The shape of the input array.</li> <li>sharded_dim (<a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)">int</a> | None) – Dimension the input array is sharded on. Defaults to 0.</li> <li>devices (Sequence[xc.Device] | None | None) – Optional sequence of devices to use. If omitted, the implicit device order used by pmap is used, which is the order of <a class="reference internal" href="_autosummary/jax.local_devices.html#jax.local_devices" title="jax.local_devices"><code class="xref py py-func docutils literal notranslate">jax.local_devices()</code></a>.</li> </ul> </dd> <dt class="field-even">Return type:</dt> <dd class="field-even"><a class="reference internal" href="#jax.sharding.PmapSharding" title="jax.sharding.PmapSharding">PmapSharding</a> </dd> </dl> </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.device_set"> property device_set: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#set" title="(in Python v3.13)">set</a>[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>]<a class="headerlink" href="#jax.sharding.PmapSharding.device_set" title="Link to this definition">#</a></dt> <dd>The set of devices that this <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> spans. In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.devices"> property devices<a class="headerlink" href="#jax.sharding.PmapSharding.devices" title="Link to this definition">#</a></dt> <dd>(self) -> ndarray </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.devices_indices_map"> devices_indices_map(global_shape)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L309-L311">[source]</a><a class="headerlink" href="#jax.sharding.PmapSharding.devices_indices_map" title="Link to this definition">#</a></dt> <dd>Returns a mapping from devices to the array slices each contains. The mapping includes all global devices, i.e., including non-addressable devices from other processes. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">global_shape (Shape) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even">Mapping[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>, Index] </dd> </dl> </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.is_equivalent_to"> is_equivalent_to(other, ndim)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L251-L254">[source]</a><a class="headerlink" href="#jax.sharding.PmapSharding.is_equivalent_to" title="Link to this definition">#</a></dt> <dd>Returns <code class="docutils literal notranslate">True</code> if two shardings are equivalent. Two shardings are equivalent if they place the same logical array shards on the same devices. For example, a <a class="reference internal" href="#jax.sharding.NamedSharding" title="jax.sharding.NamedSharding"><code class="xref py py-class docutils literal notranslate">NamedSharding</code></a> may be equivalent to a <a class="reference internal" href="#jax.sharding.PositionalSharding" title="jax.sharding.PositionalSharding"><code class="xref py py-class docutils literal notranslate">PositionalSharding</code></a> if both place the same shards of the array on the same devices. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd"><ul class="simple"> <li>self (<a class="reference internal" href="#jax.sharding.PmapSharding" title="jax.sharding.PmapSharding">PmapSharding</a>)</li> <li>other (<a class="reference internal" href="#jax.sharding.PmapSharding" title="jax.sharding.PmapSharding">PmapSharding</a>)</li> <li>ndim (<a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)">int</a>)</li> </ul> </dd> <dt class="field-even">Return type:</dt> <dd class="field-even"><a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a> </dd> </dl> </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.is_fully_addressable"> property is_fully_addressable: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="headerlink" href="#jax.sharding.PmapSharding.is_fully_addressable" title="Link to this definition">#</a></dt> <dd>Is this sharding fully addressable? A sharding is fully addressable if the current process can address all of the devices named in the <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a>. <code class="docutils literal notranslate">is_fully_addressable</code> is equivalent to “is_local” in multi-process JAX. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.is_fully_replicated"> property is_fully_replicated: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="headerlink" href="#jax.sharding.PmapSharding.is_fully_replicated" title="Link to this definition">#</a></dt> <dd>Is this sharding fully replicated? A sharding is fully replicated if each device has a complete copy of the entire data. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.memory_kind"> property memory_kind: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a> | <a class="reference external" href="https://docs.python.org/3/library/constants.html#None" title="(in Python v3.13)">None</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L316-L322">[source]</a><a class="headerlink" href="#jax.sharding.PmapSharding.memory_kind" title="Link to this definition">#</a></dt> <dd>Returns the memory kind of the sharding. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.num_devices"> property num_devices: <a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)">int</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L301-L304">[source]</a><a class="headerlink" href="#jax.sharding.PmapSharding.num_devices" title="Link to this definition">#</a></dt> <dd>Number of devices that the sharding contains. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.shard_shape"> shard_shape(global_shape)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L343-L367">[source]</a><a class="headerlink" href="#jax.sharding.PmapSharding.shard_shape" title="Link to this definition">#</a></dt> <dd>Returns the shape of the data on each device. The shard shape returned by this function is calculated from <code class="docutils literal notranslate">global_shape</code> and the properties of the sharding. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">global_shape (Shape) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even">Shape </dd> </dl> </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.sharding_spec"> property sharding_spec<a class="headerlink" href="#jax.sharding.PmapSharding.sharding_spec" title="Link to this definition">#</a></dt> <dd>(self) -> jax::ShardingSpec </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.PmapSharding.with_memory_kind"> with_memory_kind(kind)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L323-L325">[source]</a><a class="headerlink" href="#jax.sharding.PmapSharding.with_memory_kind" title="Link to this definition">#</a></dt> <dd>Returns a new Sharding instance with the specified memory kind. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">kind (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a>) </dd> </dl> </dd></dl> </dd></dl> <dl class="py class"> <dt class="sig sig-object py" id="jax.sharding.GSPMDSharding"> class jax.sharding.GSPMDSharding<a class="headerlink" href="#jax.sharding.GSPMDSharding" title="Link to this definition">#</a></dt> <dd>Bases: <a class="reference internal" href="#jax.sharding.Sharding" title="jaxlib.xla_extension.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.GSPMDSharding.device_set"> property device_set: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#set" title="(in Python v3.13)">set</a>[<a class="reference internal" href="_autosummary/jax.Device.html#jax.Device" title="jax.Device">Device</a>]<a class="headerlink" href="#jax.sharding.GSPMDSharding.device_set" title="Link to this definition">#</a></dt> <dd>The set of devices that this <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a> spans. In multi-controller JAX, the set of devices is global, i.e., includes non-addressable devices from other processes. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.GSPMDSharding.is_fully_addressable"> property is_fully_addressable: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="headerlink" href="#jax.sharding.GSPMDSharding.is_fully_addressable" title="Link to this definition">#</a></dt> <dd>Is this sharding fully addressable? A sharding is fully addressable if the current process can address all of the devices named in the <a class="reference internal" href="#jax.sharding.Sharding" title="jax.sharding.Sharding"><code class="xref py py-class docutils literal notranslate">Sharding</code></a>. <code class="docutils literal notranslate">is_fully_addressable</code> is equivalent to “is_local” in multi-process JAX. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.GSPMDSharding.is_fully_replicated"> property is_fully_replicated: <a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)">bool</a><a class="headerlink" href="#jax.sharding.GSPMDSharding.is_fully_replicated" title="Link to this definition">#</a></dt> <dd>Is this sharding fully replicated? A sharding is fully replicated if each device has a complete copy of the entire data. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.GSPMDSharding.memory_kind"> property memory_kind: <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a> | <a class="reference external" href="https://docs.python.org/3/library/constants.html#None" title="(in Python v3.13)">None</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L641-L644">[source]</a><a class="headerlink" href="#jax.sharding.GSPMDSharding.memory_kind" title="Link to this definition">#</a></dt> <dd>Returns the memory kind of the sharding. </dd></dl> <dl class="py property"> <dt class="sig sig-object py" id="jax.sharding.GSPMDSharding.num_devices"> property num_devices: <a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)">int</a><a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L633-L636">[source]</a><a class="headerlink" href="#jax.sharding.GSPMDSharding.num_devices" title="Link to this definition">#</a></dt> <dd>Number of devices that the sharding contains. </dd></dl> <dl class="py method"> <dt class="sig sig-object py" id="jax.sharding.GSPMDSharding.with_memory_kind"> with_memory_kind(kind)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/sharding_impls.py#L645-L647">[source]</a><a class="headerlink" href="#jax.sharding.GSPMDSharding.with_memory_kind" title="Link to this definition">#</a></dt> <dd>Returns a new Sharding instance with the specified memory kind. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd">kind (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a>) </dd> <dt class="field-even">Return type:</dt> <dd class="field-even"><a class="reference internal" href="#jax.sharding.GSPMDSharding" title="jax.sharding.GSPMDSharding">GSPMDSharding</a> </dd> </dl> </dd></dl> </dd></dl> <dl class="py class"> <dt class="sig sig-object py" id="jax.sharding.PartitionSpec"> class jax.sharding.PartitionSpec(*partitions)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/partition_spec.py#L35-L77">[source]</a><a class="headerlink" href="#jax.sharding.PartitionSpec" title="Link to this definition">#</a></dt> <dd>Tuple describing how to partition an array across a mesh of devices. Each element is either <code class="docutils literal notranslate">None</code>, a string, or a tuple of strings. See the documentation of <a class="reference internal" href="#jax.sharding.NamedSharding" title="jax.sharding.NamedSharding"><code class="xref py py-class docutils literal notranslate">jax.sharding.NamedSharding</code></a> for more details. This class exists so JAX’s pytree utilities can distinguish a partition specifications from tuples that should be treated as pytrees. </dd></dl> <dl class="py class"> <dt class="sig sig-object py" id="jax.sharding.Mesh"> class jax.sharding.Mesh(devices, axis_names, *, axis_types=None)<a class="reference external" href="https://github.com/jax-ml/jax/blob/main/jax/_src/mesh.py#L197-L420">[source]</a><a class="headerlink" href="#jax.sharding.Mesh" title="Link to this definition">#</a></dt> <dd>Declare the hardware resources available in the scope of this manager. In particular, all <code class="docutils literal notranslate">axis_names</code> become valid resource names inside the managed block and can be used e.g. in the <code class="docutils literal notranslate">in_axis_resources</code> argument of <a class="reference internal" href="jax.experimental.pjit.html#jax.experimental.pjit.pjit" title="jax.experimental.pjit.pjit"><code class="xref py py-func docutils literal notranslate">jax.experimental.pjit.pjit()</code></a>. Also see JAX’s multi-process programming model (<a class="reference external" href="https://jax.readthedocs.io/en/latest/multi_process.html">https://jax.readthedocs.io/en/latest/multi_process.html</a>) and the Distributed arrays and automatic parallelization tutorial (<a class="reference external" href="https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html">https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html</a>) If you are compiling in multiple threads, make sure that the <code class="docutils literal notranslate">with Mesh</code> context manager is inside the function that the threads will execute. <dl class="field-list simple"> <dt class="field-odd">Parameters:</dt> <dd class="field-odd"><ul class="simple"> <li>devices (np.ndarray) – A NumPy ndarray object containing JAX device objects (as obtained e.g. from <a class="reference internal" href="_autosummary/jax.devices.html#jax.devices" title="jax.devices"><code class="xref py py-func docutils literal notranslate">jax.devices()</code></a>).</li> <li>axis_names (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#tuple" title="(in Python v3.13)">tuple</a>[MeshAxisName, ...]) – A sequence of resource axis names to be assigned to the dimensions of the <code class="docutils literal notranslate">devices</code> argument. Its length should match the rank of <code class="docutils literal notranslate">devices</code>.</li> <li>axis_types (MeshAxisType | None)</li> </ul> </dd> </dl> Examples <div class="doctest highlight-default notranslate"><div class="highlight"><pre>>>> from jax.experimental.pjit import pjit >>> from jax.sharding import Mesh >>> from jax.sharding import PartitionSpec as P >>> import numpy as np ... >>> inp = np.arange(16).reshape((8, 2)) >>> devices = np.array(jax.devices()).reshape(4, 2) ... >>> # Declare a 2D mesh with axes `x` and `y`. >>> global_mesh = Mesh(devices, ('x', 'y')) >>> # Use the mesh object directly as a context manager. >>> with global_mesh: ... out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp) </pre></div> </div> <div class="doctest highlight-default notranslate"><div class="highlight"><pre>>>> # Initialize the Mesh and use the mesh as the context manager. >>> with Mesh(devices, ('x', 'y')) as global_mesh: ... out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp) </pre></div> </div> <div class="doctest highlight-default notranslate"><div class="highlight"><pre>>>> # Also you can use it as `with ... as ...`. >>> global_mesh = Mesh(devices, ('x', 'y')) >>> with global_mesh as m: ... out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp) </pre></div> </div> <div class="doctest highlight-default notranslate"><div class="highlight"><pre>>>> # You can also use it as `with Mesh(...)`. >>> with Mesh(devices, ('x', 'y')): ... out = pjit(lambda x: x, in_shardings=None, out_shardings=None)(inp) </pre></div> </div> </dd></dl> </section> </section> </article> <footer class="prev-next-footer"> <div class="prev-next-area"> <a class="left-prev" href="_autosummary/jax.random.weibull_min.html" title="previous page"> <div class="prev-next-info"> previous jax.random.weibull_min </div> </a> <a class="right-next" href="jax.debug.html" title="next page"> <div class="prev-next-info"> next <code class="docutils literal notranslate">jax.debug</code> module </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="#classes">Classes</a><ul class="visible nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding"><code class="docutils literal notranslate">Sharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.addressable_devices"><code class="docutils literal notranslate">Sharding.addressable_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.addressable_devices_indices_map"><code class="docutils literal notranslate">Sharding.addressable_devices_indices_map()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.device_set"><code class="docutils literal notranslate">Sharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.devices_indices_map"><code class="docutils literal notranslate">Sharding.devices_indices_map()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.is_equivalent_to"><code class="docutils literal notranslate">Sharding.is_equivalent_to()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.is_fully_addressable"><code class="docutils literal notranslate">Sharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.is_fully_replicated"><code class="docutils literal notranslate">Sharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.memory_kind"><code class="docutils literal notranslate">Sharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.num_devices"><code class="docutils literal notranslate">Sharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.shard_shape"><code class="docutils literal notranslate">Sharding.shard_shape()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Sharding.with_memory_kind"><code class="docutils literal notranslate">Sharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding"><code class="docutils literal notranslate">SingleDeviceSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.device_set"><code class="docutils literal notranslate">SingleDeviceSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.devices_indices_map"><code class="docutils literal notranslate">SingleDeviceSharding.devices_indices_map()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.is_fully_addressable"><code class="docutils literal notranslate">SingleDeviceSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.is_fully_replicated"><code class="docutils literal notranslate">SingleDeviceSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.memory_kind"><code class="docutils literal notranslate">SingleDeviceSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.num_devices"><code class="docutils literal notranslate">SingleDeviceSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.SingleDeviceSharding.with_memory_kind"><code class="docutils literal notranslate">SingleDeviceSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding"><code class="docutils literal notranslate">NamedSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.addressable_devices"><code class="docutils literal notranslate">NamedSharding.addressable_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.device_set"><code class="docutils literal notranslate">NamedSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.is_fully_addressable"><code class="docutils literal notranslate">NamedSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.is_fully_replicated"><code class="docutils literal notranslate">NamedSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.memory_kind"><code class="docutils literal notranslate">NamedSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.mesh"><code class="docutils literal notranslate">NamedSharding.mesh</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.num_devices"><code class="docutils literal notranslate">NamedSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.spec"><code class="docutils literal notranslate">NamedSharding.spec</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.NamedSharding.with_memory_kind"><code class="docutils literal notranslate">NamedSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding"><code class="docutils literal notranslate">PositionalSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.device_set"><code class="docutils literal notranslate">PositionalSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.is_fully_addressable"><code class="docutils literal notranslate">PositionalSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.is_fully_replicated"><code class="docutils literal notranslate">PositionalSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.memory_kind"><code class="docutils literal notranslate">PositionalSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.num_devices"><code class="docutils literal notranslate">PositionalSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PositionalSharding.with_memory_kind"><code class="docutils literal notranslate">PositionalSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding"><code class="docutils literal notranslate">PmapSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.default"><code class="docutils literal notranslate">PmapSharding.default()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.device_set"><code class="docutils literal notranslate">PmapSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.devices"><code class="docutils literal notranslate">PmapSharding.devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.devices_indices_map"><code class="docutils literal notranslate">PmapSharding.devices_indices_map()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.is_equivalent_to"><code class="docutils literal notranslate">PmapSharding.is_equivalent_to()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.is_fully_addressable"><code class="docutils literal notranslate">PmapSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.is_fully_replicated"><code class="docutils literal notranslate">PmapSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.memory_kind"><code class="docutils literal notranslate">PmapSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.num_devices"><code class="docutils literal notranslate">PmapSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.shard_shape"><code class="docutils literal notranslate">PmapSharding.shard_shape()</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.sharding_spec"><code class="docutils literal notranslate">PmapSharding.sharding_spec</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PmapSharding.with_memory_kind"><code class="docutils literal notranslate">PmapSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding"><code class="docutils literal notranslate">GSPMDSharding</code></a><ul class="nav section-nav flex-column"> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.device_set"><code class="docutils literal notranslate">GSPMDSharding.device_set</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.is_fully_addressable"><code class="docutils literal notranslate">GSPMDSharding.is_fully_addressable</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.is_fully_replicated"><code class="docutils literal notranslate">GSPMDSharding.is_fully_replicated</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.memory_kind"><code class="docutils literal notranslate">GSPMDSharding.memory_kind</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.num_devices"><code class="docutils literal notranslate">GSPMDSharding.num_devices</code></a></li> <li class="toc-h4 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.GSPMDSharding.with_memory_kind"><code class="docutils literal notranslate">GSPMDSharding.with_memory_kind()</code></a></li> </ul> </li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.PartitionSpec"><code class="docutils literal notranslate">PartitionSpec</code></a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#jax.sharding.Mesh"><code class="docutils literal notranslate">Mesh</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"> By The JAX authors </div> <div class="footer-item"> © Copyright 2024, The JAX Authors. </div> <div class="footer-item"> </div> <div class="footer-item"> </div> </div> </footer> </main> </div> </div>  <script src="_static/scripts/bootstrap.js?digest=5b4479735964841361fd"></script> <script src="_static/scripts/pydata-sphinx-theme.js?digest=5b4479735964841361fd"></script> <footer class="bd-footer"> </footer> </body> </html>