4675

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 »
Public API: jax package »
jax.dlpack module
Edit on GitHub

Warning

This page was created from a pull request.

jax.dlpack module¶

Next Previous

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

Built with Sphinx using a theme provided by Read the Docs.