jax

jax-mlApache-2.0活跃维护

把 NumPy 风格的 Python 代码自动微分、向量化并 JIT 编译到 GPU/TPU，写起来像 NumPy，跑起来像高性能计算。

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

jax-v0.12.0.1发布于 2026-10-22 05:37 UTC36,602star +19 / 7 天3,998fork同步于 2026-11-10 17:20 UTC

自动微分 JIT编译 GPU加速 NumPy 函数式变换

对比 GitHub 官网 jax

怎么上手

由 README 提炼，执行前请确认

原文

pip install -U jax
pip install -U "jax[cuda13]"

硬件要求按 README 说明

CPU 在 Linux x86_64/aarch64、Mac aarch64、Windows x86_64 均支持；NVIDIA GPU 支持 Linux x86_64/aarch64，Windows 仅 WSL2 实验性；Google TPU 仅 Linux x86_64；AMD GPU 支持 Linux x86_64（Windows WSL2 实验性）；Apple GPU 在 Mac aarch64 为实验性；Intel GPU 为实验性。README 未给出具体显存要求。

需要 CUDA支持 ROCm / AMD

下载：源码 zip·GitHub Releases（安装包）本站不镜像安装包

项目说明

由 README 自动翻译提炼

没人工点评？写下你的经验

用 XLA 编译并变换 NumPy 程序的库，让自动微分、向量化和 JIT 能任意组合

把 NumPy 风格的 Python 代码自动微分、向量化并 JIT 编译到 GPU/TPU，写起来像 NumPy，跑起来像高性能计算。

核心能力

用 jax.grad 做反向与正向自动微分，可任意阶求导
用 jax.jit 通过 XLA 把纯函数编译到 GPU/TPU 等加速器
用 jax.vmap 自动向量化，把循环下沉到算子层
用 make_mesh 与 PartitionSpec 做显式分片和自动并行
三种并行模式：自动、显式分片、手动 per-device 编程

容易踩的坑

GPU 支持需按平台安装对应扩展包（如 jax[cuda13]、jax[tpu]、jax[rocm7-local]），只装 jax 默认是 CPU 版；jax.jit 会限制函数内可用的 Python 控制流，需参考控制流教程；README 自称研究项目，存在 sharp edges，建议先读 Gotchas Notebook。

内容由模型根据 README 生成，可能不完整或不准确，以官方文档为准。看原文 README

最近发版

5 个版本jax-v0.12.0.1

全部 release

jax-v0.12.0.1JAX v0.12.0.1最新正式版2026-10-22 · 2026-10-22 05:37 UTC
- Bug fixes
  - Built C++-generated intrinsic LLVM IR on Windows with clang-cl instead of emitting empty bitcode, enabling vectorized CPU intrinsics on Windows (https://github.com/openxla/xla/commit/3af77b5776fbc41a65d54876cc837bdc57540005).
jax-v0.12.0JAX v0.12.02026-10-19 · 2026-10-19 11:18 UTC
查看更新说明
- New features
  
  Added jax.scipy.special.erfcinv for the inverse complementary error function.
  
  Added jax.lax.polynomial for polynomial evaluation. The polynomial primitive is faster and uses less memory when computing gradients.
  
  Added jax.custom_vjp.defremat, which customizes how a jax.custom_vjp function is rematerialized when it's differentiated under jax.remat. It requires the jax_custom_vjp3 and jax_remat3 implementations.
  
  Added a remat option to jax.custom_gradient, for writing jax.custom_vjp.defremat rules with closures.
  
  arr.at[...].get() now accepts a strategy argument that allows choosing whether the operation lowers to gather, dynamic_slice, or slice.
  
  Added a generic_batching option to jax.custom_vjp.defvjp, for batching under jax.vmap when a primal output is shared across the batch. See jax-301-vmap-shared-outputs.
  
  jax.nn.get_scaled_dot_general_config accepts 'mxfp4': E2M1 elements with one E8M0 scale per block of 32, for jax.nn.scaled_dot_general. The quantization is executed natively on AMD gfx950 GPUs and by a dequantize-and-dot fallback elsewhere.
  
  JAX now registers a proxy backend, which connects to an IFRT Proxy server such as Pathways on Cloud. To use it, set JAX_PLATFORMS=proxy and set JAX_BACKEND_TARGET to the server address, e.g. grpc://localhost:29000.
  
  Added jax.sharding.no_collectives, a context manager that raises an error at trace time in explicit sharding mode if an operation or jax.sharding.reshard introduces a collective across any specified mesh axis.
  
  Added jax.custom_vjp.defremat_with_logs, allowing jax.custom_vjp.defremat rules (and jax.custom_gradient with remat=True and with_logs=True) to log values from the backward pass (#40956).
  
  jax.scipy.linalg.eigh now supports generalized Hermitian eigenproblems (b != None) for problem types type=1, 2, and 3 (#41162, #41197).
  
  Added jax.profiler.ProfileSession and jax.profiler.export_profile_data for session- and context-manager-based profiling.
  
  Added jax.experimental.layout.Layout.for_array to query the default device layout for an array.
  
  Added a strict_in_out argument to jax.experimental.program_order to place optimization barriers on the first input and/or last output of the scope.
  
  Added jax.experimental.thunky for authoring, inspecting, splicing, and executing XLA:GPU ThunkExecutor programs directly from Python.
  
  Added JVP rule for jax.lax.nextafter, enabling forward- and reverse-mode autodiff through nextafter (#38679).
  
  Widened jax.random.ball's p parameter type from float to RealArray, allowing array-valued norm orders (#41546).
  
  Added support for DLPack 1.0 (DLManagedTensorVersioned) and read-only DLPack tensors (#40789).
  
  [oneAPI GPU] Added SYCL/oneMKL kernels for jax.lax.linalg.lu_pivots_to_permutation, jax.lax.linalg.cholesky_update, jax.lax.linalg.tridiagonal_solve (with perturb_singular=True), jax.lax.linalg.tridiagonal (sytrd/hetrd), symmetric rank-$k$ updates (syrk), and scaled_matmul, and set jax_oneapi_visible_devices in jax.distributed.initialize.
  
  [ROCm] Added gfx1250 to the default target architectures in ROCm builds (#40877).
- Breaking changes
  
  Removed jax.custom_remat. Use jax.custom_vjp.defremat instead: custom_remat(f, f_fwd, f_rem, f_bwd) becomes f = jax.custom_vjp(f) followed by f.defremat(f_fwd, f_rem, f_bwd).
  
  On x86-64, JAX and jaxlib now require an Intel Haswell or newer CPU (roughly 2013 onwards, requiring AVX2 and FMA support). Older CPUs that only support AVX are no longer suppo
jax-v0.11.2JAX v0.11.22026-09-17 · 2026-09-17 23:43 UTC
查看更新说明
- New features
  
  Added jax.numpy.minmax (and jnp.minmax), which returns both the minimum and maximum of an array, matching NumPy 2.3+ (#40089).
  
  Added jax.lax.log2 and primitive jax.lax.log2_p, making log2 a first-class primitive in JAX (jax.numpy.log2 now lowers via jax.lax.log2).
  
  Added jax.lax.one_minus_square primitive to accurately compute 1 - x^2 near $\pm 1$ and provide accurate derivatives near $0$.
  
  Added jax.export.symbolic_dim_bounds for querying conservative bounds on symbolic dimension expressions (#40006).
  
  Added frozendict support to JAX pytrees for Python 3.15 (PEP 814).
  
  jax.distributed.initialize can now secure the coordination service with mutual TLS via the new mtls_cert_file, mtls_key_file, mtls_ca_file, mtls_peer_uri_prefix and verify_secure_credentials arguments (or the JAX_MTLS_CERT_FILE, JAX_MTLS_KEY_FILE, JAX_MTLS_CA_FILE, JAX_MTLS_PEER_URI_PREFIX and JAX_DISTRIBUTED_VERIFY_SECURE_CREDENTIALS environment variables).
  
  Added cluster detection support for Open MPI 5 in jax.distributed.initialize (#40512).
  
  Added support for reading TPU_PROCESS_ADDRESSES_PATH in GKE TPU clusters.
  
  Widened jax.random.generalized_normal's p parameter type from float to RealArray, allowing array-valued shape parameters (#40126).
  
  Added exclude_argnames argument to jax.experimental.program_order.
  
  [oneAPI GPU] Added oneMKL LAPACK solver kernels for QR decomposition (geqrf, orgqr/ungqr, ormqr/unmqr), LU decomposition (getrf), symmetric/Hermitian eigenvalue decomposition (syevd/heevd), SVD (gesvd), and hybrid solver kernels (geqp3, eig) (#40000, #40186, #40543).
  
  jaxlib wheels now ship C++ FFI extension headers (collectives.h, record.h) to support out-of-tree plugins (#40333).
- Breaking changes
  
  Removed deprecated jax.experimental.shard_alike. Use explicit sharding mode instead (see sharding).
- Changes
  
  Optimized large JAX jax.sharding.Mesh construction by avoiding redundant device array allocations and copies.
  
  Thread-safety improvements across JAX and jaxlib for free-threaded Python (Python 3.13t, 3.14t, 3.15t).
  
  inline=True in jax.jit now corresponds to jax.Inline.JAX_LATE instead of jax.Inline.JAX_EARLY.
  
  The minimum CuDNN version for CUDA 12 is v9.10.2.
  
  JAX now uses Bazel 8.7.0 to build from source.
  
  JAX now uses Bzlmod for its Bazel builds instead of WORKSPACE.
  
  On GPU, multi-dimensional inverse real FFTs (jax.numpy.fft.irfftn, jax.numpy.fft.irfft2 and jax.lax.fft with FftType.IRFFT) are again lowered to a single C2R transform, as before JAX 0.10.0, instead of an IFFT over the outer axes and a 1-D IRFFT with two transposes. The input is first made Hermitian-symmetric along the outer axes, which does not change the result under NumPy's convention (only the last axis is assumed symmetric), so results are unchanged while the transform is ~1.4x faster at typical sizes.
  
  Re-enabled scan-based lowering of cumulative sums (jax.numpy.cumsum) on GPU, improving performance.
  
  jax.numpy.tri now returns an array with the default float dtype when the dtype argument is not specified. Previously it always returned float32 (#40242).
  
  jax.numpy.unique with axis specified now matches NumPy's output shape for arrays that are empty along the given axis, instead of fabricating a phantom slice for fully-empty inputs.
  
  Improved numerical accuracy of jax.numpy.log2 by pre-computing the 1 / log(2) constant factor (#40430).
  
  Added out_sharding parameter to jax.numpy.histogram.
  
  Updated jax.remat's prevent_cse argument signature to accept bool | Sequence[bool], matching jax.checkpoint.
  
  Made jax.experimental.checkify error code
jax-v0.11.1JAX v0.11.12026-08-17 · 2026-08-17 20:45 UTC
查看更新说明
- New features
  
  Added an error check for trying to deserialize JAX exports that are older than the backwards compatibility window. Without this check the deserialization of expired artifacts may succeed and then result in obscure downstream errors. Added a configuration flag --jax_export_deserialize_expired_versions to temporarily bypass the error check. See https://docs.jax.dev/en/latest/export/export.html#compatibility-guarantees.
  
  Added jax.numpy.top_k, which implements numpy.top_k, added in in NumPy v2.6.0 (#39729).
- Breaking changes
  
  The exec_time_optimization_effort and memory_fitting_effort flags have been removed in favor of the EffortLevel enum.
  
  JAX does not support anymore deserialization of Exported modules from before January 15th, 2026 because they are beyond the backwards compatibility window. On that date we added support to serialize shardings as NamedSharding, and now that is the only sharding serialization that is supported.
  
  jnp.take_along_axis now always defaults wrap_negative_indices to True. It used to default to False for mode=promise_in_bounds and True otherwise. (This also means None is no longer a valid value for wrap_negative_indices.)
- Deprecations
  
  The fields in_shardings_hlo and out_shardings_hlo of jax.export.Exported have been deprecated for a while. Now accessing them raises a warning. Use in_shardings_jax and out_shardings_jax instead.
- Changes
  
  The cuDNN fused attention backward pass (used by jax.nn.dot_product_attention with implementation='cudnn') no longer computes a bias gradient when the only attention bias comes from a boolean mask, whose gradient no caller can request. Bias gradients for an explicit bias or a non-boolean mask are unchanged (#34685).
  
  jax.numpy.meshgrid, jax.numpy.ogrid, and jax.numpy.broadcast_arrays now return tuples rather than lists in order to align with NumPy>2.0 and the Array API specification. (#39783, #39789, #39802)
  
  When jax.grad or jax.value_and_grad rejects a function with a non-scalar output, the error message now suggests reducing the output to a scalar (e.g. with output.sum()), using jax.jacobian, or reshaping size-1 outputs (#2303).
  
  When indexing with non-static or traced slice indices, the error message now suggests using jax.lax.dynamic_slice, jax.lax.dynamic_update_slice, or jax.ds, and shows tracer provenance (#7222).
  
  PyTree metadata equality comparison failures now report the specific registered pytree type that caused the error (#13027).
- Bug fixes
  
  jax.numpy.linalg.det and jax.numpy.linalg.slogdet now use a closed-form LU decomposition with row pivoting for 2x2 and 3x3 matrices instead of closed-form polynomial expansions to avoid numerical instability and catastrophic cancellation (#39905).
  
  The batching rules of the cuDNN fused attention primitives (used by jax.nn.dot_product_attention with implementation='cudnn') now support operands that do not carry the vmap axis, including a shared bias or mask. Previously jax.jacobian, jax.vmap with partial in_axes, and jax.vmap of a VJP or of jax.grad failed with a reshape TypeError (#38495).
  
  jax.vmap of fp8 cuDNN fused attention now works: its batching rules additionally mislabeled or dropped the amax outputs and restored output shapes incorrectly, so previously no vmap of the fp8 path succeeded at all. The amax outputs are whole-batch statistics and do not carry the vmap axis; vmap over the scale/descale operands raises a clear NotImplementedError.
  
  Setting jax_compiler_enable_remat_pass to False now adds rematerialization to the set of disabled XLA passes instead of overwriting it, so HLO passes disabled via `XLA_FL
jax-v0.11.0JAX v0.11.02026-07-16 · 2026-07-16 19:53 UTC
查看更新说明
- New features
  
  Added a doc on defining custom derivative rules with the experimental hijax API (hijax-custom-derivatives), along with jax.experimental.hijax helpers for deriving VJPHiPrimitive autodiff rules from a jvp or lin rule: linearize_from_jvp with apply_derived_linearization, vjp_fwd_from_jvp with transpose_jvp, vjp_fwd_from_lin with transpose_linearized, and jvp_from_lin.
  
  Added jax.custom_remat to the top-level jax namespace, for per-function control of rematerialization under the new jax_remat3 implementation.
  
  jax.checkpoint_policies is now a submodule rather than a namespace object (so from jax.checkpoint_policies import ... now works; attribute access is unchanged), and it additionally exposes the name-based policy classes SaveOnlyTheseNames, SaveAnyNamesButThese, and SaveAndOffloadOnlyTheseNames.
  
  Added jax.Inline enum for specify inlining policies to jax.jit.
- Breaking changes
  
  The deprecated module jax.cloud_tpu_init was removed. This did nothing and references to it can be safely removed.
  
  Support for Python 3.11, NumPy 2.0, and SciPy 1.14 has been dropped, per the deprecation policy.
  
  Support for Python 3.13 free-threaded (3.13t) has been dropped. Python 3.13 free-threaded was an experimental build needed to bootstrap free threading support. Now that Python 3.14t is stable, it is time to drop the experimental build. Other parts of the Python ecosystem (e.g. cibuildwheel, scipy) are making similar moves.
  
  jax.numpy.empty and jax.numpy.empty_like now produce uninitialized arrays, similar to their NumPy counterparts. Prior to v0.11.0, they produced arrays initialized to zeros. To recover the previous behavior, use jax.numpy.zeros or jax.numpy.zeros_like instead.
- Deprecations
  
  Passing 2-dimensional arrays (or mixed 2D and 3D arrays) to jax.numpy.cross is deprecated and will be removed in JAX 0.12.0, aligning with NumPy 2.5 behavior.
  
  Several previously-deprecated APIs from jax.core have been removed, including CallPrimitive, DebugInfo, DropVar, Effect, Effects, InconclusiveDimensionOperation, JaxprTypeError, abstract_token, check_jaxpr, concrete_or_error, find_top_trace, gensym, get_opaque_trace_state, is_concrete, is_constant_dim, is_constant_shape, jaxprs_in_params, new_jaxpr_eqn, no_effects, nonempty_axis_env_DO_NOT_USE, primal_dtype_to_tangent_dtype, unsafe_am_i_under_a_jit_DO_NOT_USE, unsafe_am_i_under_a_vmap_DO_NOT_USE, unsafe_get_axis_names_DO_NOT_USE, valid_jaxtype, JaxprPpContext, JaxprPpSettings, OutputType, aval_mapping_handlers, call, concretization_function_error, custom_typechecks, literalable_types, no_axis_name, and trace_ctx.
  
  Several previously-deprecated APIs from jax.interpreters.pxla have been removed, including Index, MeshAxisName, MeshExecutable, global_aval_to_result_handler, global_result_handlers, are_hlo_shardings_equal, is_hlo_sharding_replicated, ArrayMapping, _UNSPECIFIED, array_mapping_to_axis_resources, and op_sharding_to_indices.

用户点评

登录用户的真实使用记录，可按版本和硬件筛选

共 0 条

还没有人点评。第一个说出真实体验的人最有价值。

README 原文

未翻译

在 GitHub 打开

Transformable numerical computing at scale

Transformations | Scaling | Install guide | Change logs | Reference docs

What is JAX?

JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.

JAX can automatically differentiate native Python and NumPy functions. It can differentiate through loops, branches, recursion, and closures, and it can take derivatives of derivatives of derivatives. It supports reverse-mode differentiation (a.k.a. backpropagation) via jax.grad as well as forward-mode differentiation, and the two can be composed arbitrarily to any order.

JAX uses XLA to compile and scale your NumPy programs on TPUs, GPUs, and other hardware accelerators. You can compile your own pure functions with jax.jit. Compilation and automatic differentiation can be composed arbitrarily.

Dig a little deeper, and you'll see that JAX is really an extensible system for composable function transformations at scale.

This is a research project, not an official Google product. Expect sharp edges. Please help by trying it out, reporting bugs, and letting us know what you think!

import jax
import jax.numpy as jnp

def predict(params, inputs):
  for W, b in params:
    outputs = jnp.dot(inputs, W) + b
    inputs = jnp.tanh(outputs)  # inputs to the next layer
  return outputs                # no activation on last layer

def loss(params, inputs, targets):
  preds = predict(params, inputs)
  return jnp.sum((preds - targets)**2)

grad_loss = jax.jit(jax.grad(loss))  # compiled gradient evaluation function
perex_grads = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0)))  # fast per-example grads

Transformations
Scaling
Current gotchas
Installation
Citing JAX
Reference documentation

Transformations

At its core, JAX is an extensible system for transforming numerical functions. Here are three: jax.grad, jax.jit, and jax.vmap.

Automatic differentiation with `grad`

Use jax.grad to efficiently compute reverse-mode gradients:

import jax
import jax.numpy as jnp

def tanh(x):
  y = jnp.exp(-2.0 * x)
  return (1.0 - y) / (1.0 + y)

grad_tanh = jax.grad(tanh)
print(grad_tanh(1.0))
# prints 0.4199743

You can differentiate to any order with grad:

print(jax.grad(jax.grad(jax.grad(tanh)))(1.0))
# prints 0.62162673

You're free to use differentiation with Python control flow:

def abs_val(x):
  if x > 0:
    return x
  else:
    return -x

abs_val_grad = jax.grad(abs_val)
print(abs_val_grad(1.0))   # prints 1.0
print(abs_val_grad(-1.0))  # prints -1.0 (abs_val is re-evaluated)

See the JAX Autodiff Cookbook and the reference docs on automatic differentiation for more.

Compilation with `jit`

Use XLA to compile your functions end-to-end with jit, used either as an @jit decorator or as a higher-order function.

import jax
import jax.numpy as jnp

def slow_f(x):
  # Element-wise ops see a large benefit from fusion
  return x * x + x * 2.0

x = jnp.ones((5000, 5000))
fast_f = jax.jit(slow_f)
%timeit -n10 -r3 fast_f(x)
%timeit -n10 -r3 slow_f(x)

Using jax.jit constrains the kind of Python control flow the function can use; see the tutorial on Control Flow and Logical Operators with JIT for more.

Auto-vectorization with `vmap`

vmap maps a function along array axes. But instead of just looping over function applications, it pushes the loop down onto the function’s primitive operations, e.g. turning matrix-vector multiplies into matrix-matrix multiplies for better performance.

Using vmap can save you from having to carry around batch dimensions in your code:

import jax
import jax.numpy as jnp

def l1_distance(x, y):
  assert x.ndim == y.ndim == 1  # only works on 1D inputs
  return jnp.sum(jnp.abs(x - y))

def pairwise_distances(dist1D, xs):
  return jax.vmap(jax.vmap(dist1D, (0, None)), (None, 0))(xs, xs)

xs = jax.random.normal(jax.random.key(0), (100, 3))
dists = pairwise_distances(l1_distance, xs)
dists.shape  # (100, 100)

By composing jax.vmap with jax.grad and jax.jit, we can get efficient Jacobian matrices, or per-example gradients:

per_example_grads = jax.jit(jax.vmap(jax.grad(loss), in_axes=(None, 0, 0)))

Scaling

To scale your computations across thousands of devices, you can use any composition of these:

Compiler-based automatic parallelization where you program as if using a single global machine, and the compiler chooses how to shard data and partition computation (with some user-provided constraints);
Explicit sharding and automatic partitioning where you still have a global view but data shardings are explicit in JAX types, inspectable using jax.typeof;
Manual per-device programming where you have a per-device view of data and computation, and can communicate with explicit collectives.

Mode	View?	Explicit sharding?	Explicit Collectives?
Auto	Global	❌	❌
Explicit	Global	✅	❌
Manual	Per-device	✅	✅

from jax.sharding import set_mesh, AxisType, PartitionSpec as P
mesh = jax.make_mesh((8,), ('data',), axis_types=(AxisType.Explicit,))
set_mesh(mesh)

# parameters are sharded for FSDP:
for W, b in params:
  print(f'{jax.typeof(W)}')  # f32[512@data,512]
  print(f'{jax.typeof(b)}')  # f32[512]

# shard data for batch parallelism:
inputs, targets = jax.device_put((inputs, targets), P('data'))

# evaluate gradients, automatically parallelized!
gradfun = jax.jit(jax.grad(loss))
param_grads = gradfun(params, (inputs, targets))

See the tutorial and advanced guides for more.

Gotchas and sharp bits

See the Gotchas Notebook.

Installation

Supported platforms

	Linux x86_64	Linux aarch64	Mac aarch64	Windows x86_64	Windows WSL2 x86_64
CPU	yes	yes	yes	yes	yes
NVIDIA GPU	yes	yes	n/a	no	experimental
Google TPU	yes	n/a	n/a	n/a	n/a
AMD GPU	yes	no	n/a	no	experimental
Apple GPU	n/a	no	experimental	n/a	n/a
Intel GPU	experimental	n/a	n/a	no	no

Instructions

Platform	Instructions
CPU	`pip install -U jax`
NVIDIA GPU	`pip install -U "jax[cuda13]"`
Google TPU	`pip install -U "jax[tpu]"`
AMD GPU (Linux, ROCm from pip)	`pip install -U "jax[rocm7-device-gfx942]" --extra-index-url https://repo.amd.com/rocm/whl-multi-arch/`
AMD GPU (Linux, system ROCm)	`pip install -U "jax[rocm7-local]"`
Intel GPU	`pip install -U "jax[oneapi]"`

See the documentation for information on alternative installation strategies. These include compiling from source, installing with Docker, using other versions of CUDA, a community-supported conda build, and answers to some frequently-asked questions.

Citing JAX

To cite this repository:

@software{jax2018github,
  author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Yash Katariya and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang},
  title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs},
  url = {http://github.com/jax-ml/jax},
  version = {0.3.13},
  year = {2018},
}

In the above bibtex entry, names are in alphabetical order, the version number is intended to be that from jax/version.py, and the year corresponds to the project's open-source release.

A nascent version of JAX, supporting only automatic differentiation and compilation to XLA, was described in a paper that appeared at SysML 2018. We're currently working on covering JAX's ideas and capabilities in a more comprehensive and up-to-date paper.

Reference documentation

For details about the JAX API, see the reference documentation.

For getting started as a JAX developer, see the developer documentation.