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

Runtime value debugging in JAX

Reference Documentation

JAX Frequently Asked Questions (FAQ)
Asynchronous dispatch
Ahead-of-time lowering and compilation
Understanding Jaxprs
Convolutions in JAX
Pytrees
jax.Array migration
Type promotion semantics
JAX Errors
Transfer guard
JAX Glossary of Terms
Change log

Advanced JAX Tutorials

The Autodiff Cookbook
Using JAX in multi-host and multi-process environments
Parallelism with JAX
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

Developer documentation

Contributing to JAX
Building from source
Internal APIs
Autodidax: JAX core from scratch
JAX Enhancement Proposals (JEPs)

API documentation

Public API: jax package

Notes

API compatibility
Python and NumPy version support policy
Concurrency
GPU memory allocation
Profiling JAX programs
Device Memory Profiling
Rank promotion warning

Theme by the Executable Book Project

.rst

Contents

core

Internal APIs

Contents

core

Internal APIs#

core#

Jaxpr(constvars, invars, outvars, eqns[, ...])

param constvars

ClosedJaxpr(jaxpr, consts)

param jaxpr

previous

Building from source

next

jax.core.Jaxpr

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