DOCSEARCHbyAlgolia
Submit Docs

JAX

repository·main·Indexed 12 days ago

https://github.com/jax-ml/jax

A high-performance Python library for numerical computing and large-scale machine learning. JAX features composable function transformations including automatic differentiation (jax.grad), XLA compilation (jax.jit), and auto-vectorization (vmap). It supports scaling across accelerators like GPUs and TPUs through various modes (Auto, Explicit, Manual) and provides specialized tools like Stax for functional neural network building and Refs for mutable array operations.

Tokens
301.6K
Snippets
759
Records
940
Agent score
89%

What's inside JAX

  1. Overview of JAX core capabilities

    main

    JAX is a high-performance array computing library designed for numerical computing and machine learning. Its core value proposition includes:

    • Familiar API: Provides a NumPy-style API to ease adoption for researchers and engineers.
    • Transformations: Supports composable function transformations, including:
      • Compilation (e.g., jit for Just-In-Time compilation)
      • Batching (e.g., vmap for auto-vectorization)
      • Automatic differentiation (e.g., grad)
      • Parallelization
    • Multi-backend execution: The same code can execute across multiple hardware backends, including CPU, GPU, and TPU.

  2. What is JAX?

    main

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

    Key capabilities include:

    • Automatic Differentiation: Differentiates native Python and NumPy functions, including through loops, branches, recursion, and closures. It supports both reverse-mode (jax.grad) and forward-mode differentiation.
    • XLA Compilation: Uses XLA to compile and scale NumPy programs on TPUs, GPUs, and other hardware accelerators via jax.jit.
    • Composable Transformations: JAX is built as an extensible system where transformations like grad, jit, and vmap can be composed arbitrarily.

  3. What is Colocated Python?

    main

    Colocated Python is an experimental API that provides a uniform way to run Python code on the hosts associated with a set of JAX devices.

    • Local Devices: If JAX devices are local, Python code runs on the local host.
    • Remote Devices: If JAX devices are remote, Python code is shipped to run on the host of those remote devices.

    This is designed to help build portable multi-host ML systems that work across both single-controller and multi-controller JAX environments.

    Warning: This is an experimental API. Its functionality and interface are subject to change without following the standard JAX compatibility policy.

  4. What is omnistaging and why is it used?

    main

    Omnistaging is a JAX core upgrade that stages out more computation from Python to XLA. Instead of staging operations based only on data dependence (whether they depend on a function argument), omnistaging stages out all jax.numpy calls within the dynamic context of a jit, pmap, or control flow primitive.

    Benefits:

    • Improved Memory Performance: Reduces memory fragmentation during tracing and produces fewer large compile-time constants for XLA.
    • Faster Tracing: Eliminates op-by-op execution at Python tracing time.
    • Simplified Internals: Simplifies JAX core and fixes various bugs.

    Mental Model Change: Instead of treating jax.numpy as a direct drop-in replacement for numpy for all tasks, think of using jax.numpy operations specifically when you want a computation to be performed on an accelerator (like a GPU). For trace-time constants or shape computations, use standard numpy.

    # Before omnistaging, this 'add' might not be staged out to XLA:
@jit
def f(x):
  y = jnp.add(1, 1)
  return x * y

# After omnistaging, the 'add' is staged out and part of the XLA HLO:
# ENTRY jit_f.8 {
#   ... 
#   add.5 = s32[] add(constant.3, constant.4)
#   multiply.6 = s32[] multiply(parameter.1, add.5)
#   ... 
# }

  5. Use jax.typing for permissive input annotations

    main

    To improve code readability and compatibility with static type checkers, JAX is introducing a jax.typing module. When writing functions that accept JAX-compatible data, use these permissive types for inputs:

    • ArrayLike: A union of anything that can be implicitly converted into an array (e.g., JAX arrays, NumPy arrays, JAX tracers, or Python/NumPy scalars).
    • DTypeLike: A union of anything convertible into a dtype (e.g., NumPy dtypes, strings, or built-in types like float).
    • ShapeLike: A union of anything convertible into a shape (e.g., sequences of integers).

    Note: These are designed to be simpler than NumPy's equivalents, as JAX does not support certain complex structures like structured dtypes or list/tuple inputs in place of arrays.

  6. Use `ClusterBarrier` for cross-block synchronization

    main

    A plgpu.ClusterBarrier is used to synchronize across block clusters rather than threads within a single block. This is required when blocks in a cluster collaborate on shared resources like SMEM or TMEM.

    Common Use Cases:

    • SMEM Reuse: Ensuring all blocks in a cluster have finished reading from a shared SMEM region before one block overwrites it with a new collective async copy.
    • TMEM Reuse (Blackwell): Ensuring all blocks have completed their reads from TMEM before the resource is reused for a subsequent collective MMA.
    def collective_smem_reuse(x_gmem, x_gmem2, y_gmem, x_smem, local_barrier, cluster_barrier):
  # Step 1: Collective copy into SMEM
  plgpu.copy_gmem_to_smem(x_gmem, x_smem, local_barrier, collective_axes="cluster")
  plgpu.barrier_wait(local_barrier)  # Wait for local copy to finish
  
  y_gmem[0] = x_smem[...]
  
  # Step 2: Synchronize cluster to ensure all blocks finished using x_smem
  plgpu.barrier_arrive(cluster_barrier)
  plgpu.barrier_wait(cluster_barrier)
  
  # Step 3: Safe to overwrite x_smem
  plgpu.copy_gmem_to_smem(x_gmem2, x_smem, local_barrier, collective_axes="cluster")
  plgpu.barrier_wait(local_barrier)
  y_gmem[1] = x_smem[...]

  7. Shape polymorphic export

    main
    When exporting, you can use dimension variables for certain input dimensions. This allows the resulting exported artifact to be used with multiple different combinations of input shapes, rather than being locked to the specific shapes used during tracing.

  8. Limitations of `StatefulPRNG`

    main

    While convenient, the stateful PRNG has several technical limitations due to its reliance on JAX's effect system (mutable refs):

    1. Sequential Dependence: The internal counter creates a sequential dependency. The compiler cannot reorder operations that depend on these random values, which may impact performance compared to batched stateless operations.
    2. Refactoring Sensitivity: Changing the order of code (e.g., adding a layer in a neural network) will change the random draws for all subsequent operations because the counter increments sequentially.
    3. Incompatibility with remat: Stateful keys cannot be used within jax.remat (rematerialization) because remat does not support mutable refs. Attempting this will result in an explicit error.
    4. No Return Values: You cannot return a StatefulPRNG object as a return value from a transformed JAX function (e.g., inside jit).

  9. Composing `jax.jit` and `jax.vmap`

    main

    JAX transformations are composable. You can wrap a vmap-ed function with jax.jit to gain compilation benefits, or wrap a jit-ed function with jax.vmap. This allows you to combine the efficiency of vectorized operations with the speed of JIT compilation.

    import jax

# Assuming auto_batch_convolve is a vmapped function
jitted_batch_convolve = jax.jit(auto_batch_convolve)

# This will run the vectorized function with JIT compilation
jitted_batch_convolve(xs, ws)

  10. Understand the `vmap` and custom differentiation semantics problem

    main

    A known issue in older JAX versions (related to the deprecated custom_transforms API) is that non-differentiation transformations like vmap or mask can effectively 'remove' custom differentiation rules.

    This happens because vmap operates by inlining/rewriting the function definition. If a function f is wrapped in a custom rule via a primitive, vmap(f) rewrites the function into its underlying constituent primitives, bypassing the custom rule associated with the original function f.

    Consequence: Applying vmap to a function with a custom derivative rule may result in the gradient being calculated using the original function's implementation rather than your custom rule. This violates the expected semantic that vmap(f)(xs) == [f(x) for x in xs].

  11. How CuTe Tensors and Layouts work

    main

    In the CuTe DSL, a Tensor is a combination of a pointer to GPU memory and a Layout. The Layout describes how to navigate the memory using coordinates.

    Instead of manually computing memory offsets, you index a tensor using its logical coordinates (e.g., a[None, tidx, bidx]), and the CuTe layout handles the mapping to the flat memory address. This allows for Layout composition: you can reshape data (e.g., from 1-D to 3-D) by combining the original layout with a new one. This operation is purely algebraic and involves no data movement or copying, making reshaping 'free'.

  12. Implement generic container types (Tuples) with hijax

    main

    You can define container types (like HiTup) where the type is parameterized by its component types. This allows for nested structures and heterogeneous elements.

    Key features:

    • Recursive Lowering: lo_ty and lower_val delegate to the component types, allowing for nested tuples or tuples containing other HiTypes.
    • Per-component Mapping: When using vmap, you can use a TupSpec to specify different in_axes and out_axes for each element in the tuple.
    • Static Parameters: Primitives like GetTupElt use static parameters (e.g., idx) passed via the params dictionary to allow indexing into the container during tracing.

    Example of a TupSpec for vmap:

    # Map the first element on axis 0, but leave the second element unbatched
out = jax.vmap(swap, in_axes=TupSpec((0, None)), out_axes=TupSpec((None, 0)), axis_size=3)(tup)
    @dataclass(frozen=True)
class TupSpec(MappingSpec):
  val: tuple  # one axis entry per component