Logo
latest

Getting Started

  • JAX Quickstart
  • How to Think in JAX
  • 🔪 JAX - The Sharp Bits 🔪
  • Tutorial: JAX 101
    • JAX As Accelerated NumPy
    • Just In Time Compilation with JAX
    • Automatic Vectorization in JAX
    • Advanced Automatic Differentiation in JAX
    • Pseudo Random Numbers in JAX
    • Working with Pytrees
    • Parallel Evaluation in JAX
    • Stateful Computations in JAX
    • Introduction to pjit

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
  • XLA in Python
  • Named axes and easy-to-revise parallelism
  • Using JAX in multi-host and multi-process environments

Notes

  • 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

Developer documentation

  • Contributing to JAX
  • Building from source
  • Running the tests
  • Type checking
  • Update documentation
  • Internal APIs
  • Autodidax: JAX core from scratch

API documentation

  • Public API: jax package
JAX
  • Docs »
  • Tutorial: JAX 101
  • Edit on GitHub

Tutorial: JAX 101¶

This is a tutorial developed by engineers and researchers at DeepMind.

Tutorials

  • JAX As Accelerated NumPy
  • Just In Time Compilation with JAX
  • Automatic Vectorization in JAX
  • Advanced Automatic Differentiation in JAX
  • Pseudo Random Numbers in JAX
  • Working with Pytrees
  • Parallel Evaluation in JAX
  • Stateful Computations in JAX
  • Introduction to pjit
Next Previous

© Copyright 2020, Google LLC. NumPy and SciPy documentation are copyright the respective authors. Revision 9a5cf7c8.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
Versions
latest
stable
4510-2
test-docs
Downloads
html
On Read the Docs
Project Home
Builds
Free document hosting provided by Read the Docs.