jaxlib

Warnings

• breaking The `jax.pmap` function is now in maintenance mode, and its default implementation has changed. Users are strongly encouraged to migrate new code to `jax.shard_map` for data parallelism.
  Fix: Rewrite `jax.pmap` usage to `jax.shard_map`. Consult the JAX migration guide for `pmap`.
• breaking The minimum supported NumPy version is now 2.0, and consequently, the minimum supported SciPy version is 1.13. Using older versions will lead to errors.
  Fix: Upgrade NumPy to version 2.0 or newer and SciPy to 1.13 or newer (`pip install --upgrade numpy scipy`).
• breaking `jax.dlpack.from_dlpack` no longer accepts a raw DLPack capsule directly. It must now be called with an array implementing the `__dlpack__` and `__dlpack_device__` protocols.
  Fix: Ensure that the input to `jax.dlpack.from_dlpack` is an object that implements the DLPack Python protocol.
• gotcha JAX arrays are immutable, unlike NumPy arrays. In-place modification operations common in NumPy (e.g., `arr[0] = 5`) are not supported and will raise an error or require explicit functional updates like `.at[idx].set(value)`.
  Fix: Adopt a functional programming style for array manipulations. Use `array = array.at[index].set(value)` for updates, which returns a new array with the modification.
• gotcha Installing `jaxlib` for NVIDIA GPUs or TPUs requires specific installation commands (e.g., `jax[cudaXX_pip]`) and often requires matching CUDA/cuDNN versions. Using `pip install jaxlib` alone will typically install a CPU-only version, and mismatched versions can lead to runtime errors or devices not being detected.
  Fix: Always refer to the official JAX installation guide for the correct, platform-specific commands. Ensure your CUDA/cuDNN versions are compatible with the `jaxlib` wheel you are installing.
• deprecated The semi-private type `jax._src.literals.TypedNdArray` is now a subclass of `np.ndarray`, not just a duck type. This change may affect code relying on `isinstance(x, np.ndarray)` or similar type checks for JAX internal types if they were previously treated as distinct from `np.ndarray`.
  Fix: If your code relies on distinguishing `TypedNdArray` from `np.ndarray`, adjust type checks accordingly. Convert to classic NumPy arrays using `np.asarray(x)` if necessary.

Install

• pip install --upgrade pip pip install --upgrade jax jaxlib For CPU-only (Linux/macOS/Windows)
• pip install --upgrade pip pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html For NVIDIA GPU (CUDA 12, pip-installed CUDA/cuDNN)
• # For specific CUDA versions or locally installed CUDA, refer to official JAX docs: # https://jax.readthedocs.io/en/latest/installation.html Other accelerator installations

Imports

• jax

  import jax

  Users typically import `jax` (the Python frontend) to interact with the JAX ecosystem, which implicitly uses `jaxlib` as its backend. Direct imports from `jaxlib` for general user code are rare.
• jax.numpy

  import jax.numpy as jnp

  JAX's NumPy-like API (`jax.numpy`) is the primary way users interact with array operations that are compiled and executed by `jaxlib`.

Quickstart

This quickstart demonstrates a basic JAX program that implicitly leverages `jaxlib` for Just-In-Time (JIT) compilation and execution on available accelerators (CPU, GPU, or TPU). It defines a simple numerical function, compiles it with `jax.jit`, and performs an operation on a JAX array. The output shows detected devices and a sample of the computation.

import jax
import jax.numpy as jnp

def my_function(x):
    return jnp.sin(x) * jnp.cos(x)

# JIT-compile the function for performance
compiled_function = jax.jit(my_function)

# Create a JAX array
x = jnp.linspace(0, 10, 1000)

# Run the compiled function
y = compiled_function(x)

print(f"JAX detected devices: {jax.devices()}")
print(f"Result array shape: {y.shape}")
print(f"First 5 elements of y: {y[:5]}")