Tools · SPBoucher
Menu

JAX

by Google (Chrome DevTools team) · Library · also known as jax-ml/jax, google/jax

Python library for NumPy-style array computing with composable transformations (jit, grad, vmap, shard_map) compiled through XLA to CPU, GPU and TPU

Apache-2.0 licenseOpen sourceFreeLocalSelf-hostedCloudSelf-hostable
Visit official site Compare

Link verified Nov 2, 2026

Overview

JAX is an Apache-2.0 Python library from Google for high-performance numerical computing and machine learning research. It exposes a NumPy-compatible API and a set of composable function transformations: jax.jit compiles functions through XLA, jax.grad differentiates them, jax.vmap vectorizes over a batch dimension and shard_map expresses per-device programs for multi-device execution. The same code runs on CPU, GPU and TPU backends, which is why JAX underlies much of Google's large-scale training and a growing number of academic codebases. The repository moved from google/jax to jax-ml/jax.

JAX itself is deliberately minimal; neural-network layers come from Flax, Equinox or Keras, optimizers from Optax and checkpointing from Orbax. Its functional style (pure functions, explicit random keys, immutable arrays) rewards careful program design and makes transformations predictable, but it is a different mental model from PyTorch's eager, object-oriented approach. Practical limits include longer compile times for large programs, a smaller ecosystem of pretrained models than PyTorch, and GPU support that depends on CUDA builds while Apple Silicon acceleration is experimental.

Capabilities

  • NumPy-compatible array API
  • Composable jit, grad, vmap and shard_map transformations
  • XLA compilation for CPU, GPU and TPU
  • Multi-device and multi-host parallelism
  • Ecosystem: Flax, Equinox, Optax, Orbax

Typical workflows

  1. Write a differentiable simulation and vectorize it over parameters
  2. Train a Flax model on TPUs with shard_map
  3. Port a NumPy algorithm to GPU with minimal changes

Strengths

  • Clean functional transformations
  • First-class TPU support
  • Apache-2.0 license

Limitations

  • Compile latency for large programs
  • Smaller pretrained-model ecosystem than PyTorch
  • Apple Silicon GPU support is experimental

Prerequisites

  • Python 3.10+; CUDA for NVIDIA GPUs or TPU access

Who it is for

  • ML and scientific-computing researchers
  • Teams training on TPUs

Alternatives & comparison

  • Metal

    Apple's low-level GPU programming framework for graphics and parallel compute on iOS, macOS and visionOS, the layer beneath MLX, Core ML and most Apple Silicon ML runtimes

    ProprietaryFree
  • MLX

    Apple's open-source array framework for machine learning on Apple Silicon, with NumPy-like Python, C++ and Swift APIs and unified CPU/GPU memory

    Open sourceOpen source
  • NumPy

    Foundational Python library for n-dimensional arrays, vectorized math, linear algebra, random numbers and the array API used by the scientific stack

    Open sourceOpen source
  • PyTorch

    Open-source deep-learning framework with eager tensors, autograd, torch.compile and distributed training, governed by the PyTorch Foundation and the base of most LLM tooling

    Open sourceOpen source
  • scikit-learn

    BSD-licensed Python machine-learning library for classification, regression, clustering, dimensionality reduction, preprocessing and model selection on NumPy and SciPy

    Open sourceFree
Compare JAX with its alternatives →

Sources & verification

Primary sources consulted for this record. Claims are tied to the fields each source supports; the automated link checker re-verifies every URL on a schedule.

Confidence
High
Reviewed
Nov 1, 2026
Links
Verified · Nov 2, 2026

Spotted an error?

Broken link, outdated license, wrong category, or a tool that no longer exists? Tell us; reports are reviewed by the editor.

MLX LM

Apple Inc. · Apple Silicon ML & MLX

In-depth entry

Python package and CLI for running, serving, quantizing and fine-tuning large language models on Apple Silicon with MLX

MITOpen sourceFreeLocalNov 2, 2026

Ray

Ray project (Anyscale and contributors) · Distributed & parallel computing

In-depth entry

Open-source Python framework for distributing tasks, actors and ML workloads across a cluster, with libraries for data processing, training, tuning, serving and reinforcement learning

Apache-2.0Open sourceFreeLocalSelf-hostedBroken

Hugging Face Hub

Hugging Face · Training & fine-tuning

In-depth entry

Hosting platform for open models, datasets and Spaces demo apps built on git repositories and Xet storage, with Inference Providers, Endpoints, Jobs and an open-source Python client and CLI

MixedFreemiumPaidLocalNov 2, 2026

scikit-learn

scikit-learn developers · Mathematical & statistical computing

BSD-licensed Python machine-learning library for classification, regression, clustering, dimensionality reduction, preprocessing and model selection on NumPy and SciPy

BSD-3-ClauseFreeOpen sourceLocalNov 2, 2026

Sentence Transformers

Hugging Face · Embeddings, reranking & retrieval

Python library for computing and training embedding, cross-encoder reranker, sparse and multi-vector retrieval models, maintained by Hugging Face with thousands of pretrained models on the Hub

Apache-2.0Open sourceFreeLocalSelf-hostedNov 2, 2026

Unsloth

Unsloth AI · Training & fine-tuning

Open-source framework for fine-tuning and running LLMs with reduced memory, covering LoRA, QLoRA, full fine-tuning and RL, GGUF export, Dynamic quantizations and a local Studio/Desktop UI

Apache-2.0Open sourceFreeLocalSelf-hostedNov 2, 2026