Logo
latest

Tutorials

  • JAX Quickstart
  • The Autodiff Cookbook
  • Autobatching log-densities example
  • Training a Simple Neural Network, with Tensorflow Datasets Data Loading

Advanced JAX Tutorials

  • 🔪 JAX - The Sharp Bits 🔪
  • Custom derivative rules for JAX-transformable Python functions
  • How JAX primitives work
  • Writing custom Jaxpr interpreters in JAX

Notes

  • Change Log
  • JAX Frequently Asked Questions (FAQ)
  • Understanding Jaxprs
  • Asynchronous dispatch
  • Concurrency
  • GPU memory allocation
  • Profiling JAX programs
  • Device Memory Profiling
  • Pytrees
  • Rank promotion warning
  • Type promotion semantics

Developer documentation

  • Building from source
  • Running the tests
  • Type checking
  • Update documentation
  • Internal APIs

API documentation

  • Public API: jax package
JAX
  • Docs »
  • Python Module Index

Python Module Index

j
 
j
- jax
    jax.core
    jax.dlpack
    jax.experimental
    jax.experimental.host_callback
    jax.experimental.loops
    jax.experimental.optimizers
    jax.experimental.stax
    jax.image
    jax.lax
    jax.lax.linalg
    jax.nn
    jax.nn.initializers
    jax.numpy
    jax.numpy.fft
    jax.numpy.linalg
    jax.ops
    jax.profiler
    jax.random
    jax.scipy.linalg
    jax.scipy.ndimage
    jax.scipy.optimize
    jax.scipy.signal
    jax.scipy.sparse.linalg
    jax.scipy.special
    jax.scipy.stats.bernoulli
    jax.scipy.stats.beta
    jax.scipy.stats.cauchy
    jax.scipy.stats.dirichlet
    jax.scipy.stats.expon
    jax.scipy.stats.gamma
    jax.scipy.stats.geom
    jax.scipy.stats.laplace
    jax.scipy.stats.logistic
    jax.scipy.stats.multivariate_normal
    jax.scipy.stats.norm
    jax.scipy.stats.pareto
    jax.scipy.stats.poisson
    jax.scipy.stats.t
    jax.scipy.stats.uniform
    jax.tree_util

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

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.