Toggle navigation sidebar

Getting Started

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

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
Internal APIs
Autodidax: JAX core from scratch
Design Notes

API documentation

Public API: jax package

Theme by the Executable Book Project

Toggle navigation

jax.example_libraries package

Warning

This page was created from a pull request (#9655).

jax.example_libraries package¶

jax.example_libraries.optimizers module
jax.example_libraries.stax module

previous

jax.distributed.initialize

next

jax.example_libraries.optimizers module

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