Toggle navigation sidebar

Toggle in-page Table of Contents

Getting Started

Installing JAX
JAX Quickstart
How to Think in JAX
🔪 JAX - The Sharp Bits 🔪

Tutorial: JAX 101

Reference Documentation

JAX Frequently Asked Questions (FAQ)
Asynchronous dispatch
Understanding Jaxprs
Convolutions in JAX
Pytrees
Type promotion semantics
JAX Errors
JAX Glossary of Terms
Change log

Advanced JAX Tutorials

The Autodiff Cookbook
Autobatching log-densities example
Training a Simple Neural Network, with tensorflow/datasets Data Loading
Custom derivative rules for JAX-transformable Python functions
How JAX primitives work
Writing custom Jaxpr interpreters in JAX
Training a Simple Neural Network, with PyTorch Data Loading
Named axes and easy-to-revise parallelism
Using JAX in multi-host and multi-process environments

Notes

API compatibility
Python and NumPy version support policy
Concurrency
GPU memory allocation
Profiling JAX programs
Device Memory Profiling
Rank promotion warning
custom_vjp and nondiff_argnums update guide
Transfer guard

Developer documentation

Contributing to JAX
Building from source
Internal APIs
Autodidax: JAX core from scratch
Design Notes

API documentation

Public API: jax package

Theme by the Executable Book Project

.rst

jax.distributed module

jax.distributed module#

initialize([coordinator_address, ...])

Initialize distributed system for topology discovery.

previous

jax.dlpack.to_dlpack

next

jax.distributed.initialize

By The JAX authors
© Copyright 2020, Google LLC. NumPy and SciPy documentation are copyright the respective authors..