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
Subpackages
jax.numpy package
jax.scipy package
jax.experimental package
jax.image package
jax.lax package
jax.nn package
jax.ops package
jax.random package
jax.tree_util package
jax.dlpack module
jax.profiler module
Just-in-time compilation (
jit
)
Automatic differentiation
Vectorization (
vmap
)
Parallelization (
pmap
)
JAX
Docs
»
Public API: jax package
»
jax.dlpack module
Edit on GitHub
jax.dlpack module
¶
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
.