7757

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
XLA in Python
MAML Tutorial with JAX
Generative Modeling by Estimating Gradients of Data Distribution in JAX
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 »
Public API: jax package »
jax.profiler module »
jax.profiler.trace_function
Edit on GitHub

jax.profiler.trace_function¶

jax.profiler.trace_function(*args, **kwargs)[source]¶

Next Previous

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

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