jax-js：纯 JavaScript 中的 JAX

官网 | API 参考文档 | 兼容性表格 | Discord 社区

jax-js 是一个面向浏览器的机器学习框架。它的目标是将 JAX 风格的高性能 CPU 和 GPU 内核引入 JavaScript，从而让你能够在网页上运行数值计算应用。

npm i @jax-js/jax

在底层实现中，它会将数组操作转换为编译器中间表示，然后利用 WebAssembly 和 WebGPU 合成内核代码。

该库完全从零开始编写，没有任何外部依赖。它与 NumPy/JAX 保持高度的 API 兼容性。由于所有计算都在客户端进行，jax-js 很可能是目前最便携的 GPU 机器学习框架，因为它可以在任何支持浏览器的环境中运行。

快速入门

import { numpy as np } from "@jax-js/jax";

// 数组操作，与 JAX/NumPy 兼容。
const x = np.array([1, 2, 3]);
const y = x.mul(4); // [4, 8, 12]

在网页中使用（CDN）

在原生 JavaScript 环境中（无需打包工具），只需通过 <script type="module"> 标签导入即可。这是在空白 HTML 页面上快速上手的最简单方式。

<script type="module">
  import { numpy as np } from "https://esm.sh/@jax-js/jax";
</script>

平台支持情况

本表格基于各浏览器的最新版本。截至 2025 年底，WebGPU 已在大多数主流浏览器中得到广泛支持。

平台	CPU (Wasm)	GPU (WebGPU)	GPU (WebGL)
Chrome / Edge	✅	✅	✅
Firefox	✅	✅ - macOS 26+	✅
Safari	✅	✅ - macOS 26+	✅
iOS	✅	✅ - iOS 26+	✅
Chrome for Android	✅	✅	✅
Firefox for Android	✅	❌	✅
Node.js	✅	❌	❌
Deno	✅	✅ - 异步	❌

示例

社区使用案例：

• autoresearch-webgpu：浏览器中的自动研究
• tanh.xyz：交互式机器学习可视化
• jax-js-bayes：声明式贝叶斯建模库

jax-js 官网上的一些演示：

• 在 MNIST 数据集上训练神经网络
• 语音克隆：Kyutai Pocket TTS
• 在浏览器中为书籍生成 CLIP 嵌入
• 目标检测：DETR ResNet-50 (ONNX)
• 流体模拟（Navier-Stokes 方程）
• 浏览器 REPL
• 矩阵乘法性能基准测试
• 二维卷积性能基准测试
• 曼德勃罗集

功能对比

以下是 jax-js 与其他流行 Web ML 运行时的简要对比：

特性	jax-js	TensorFlow.js	onnxruntime-web
概述
API 风格	JAX/NumPy	类 TensorFlow	静态 ONNX 图
最新版本	2026	⚠️ 2024	2026
速度	最快	较快	最快
打包大小（gzip）	80 KB	269 KB	90 KB + 24 MB Wasm
自动微分与 JIT
梯度	✅	✅	❌
雅可比矩阵和海森矩阵	✅	❌	❌
jvp() 正向微分	✅	❌	❌
jit() 内核融合	✅	❌	❌
vmap() 自动向量化	✅	❌	❌
图捕获	✅	❌	✅
后端与数据
WebGPU 后端	✅	🟡 预览	✅
WebGL 后端	✅	✅	✅
Wasm（CPU）后端	✅	✅	✅
勇于执行的数组 API	✅	✅	❌
运行 ONNX 模型	🟡 部分	❌	✅
读取 safetensors 文件	✅	❌	❌
Float64	✅	❌	❌
Float32	✅	✅	✅
Float16	✅	❌	✅
BFloat16	❌	❌	❌
压缩 Uint8	❌	❌	🟡 部分
混合精度	✅	❌	✅
多设备混合使用	✅	❌	❌
操作与数值计算
算术函数	✅	✅	✅
矩阵乘法	✅	✅	✅
通用 einsum	✅	🟡 部分	🟡 部分
排序	✅	❌	❌
激活函数	✅	✅	✅
NaN/Inf 数值计算	✅	✅	✅
基础卷积	✅	✅	✅
n 维卷积	✅	❌	✅
步进/空洞卷积	✅	✅	✅
Cholesky 分解、最小二乘问题	✅	❌	❌
LU 分解、线性方程求解、行列式	✅	❌	❌
奇异值分解	❌	❌	❌
FFT	✅	✅	✅
基础随机数生成（均匀分布、正态分布）	✅	✅	✅
高级随机数生成	✅	❌	❌

教程

在 jax-js 中进行编程看起来与 JAX 非常相似，只是使用 JavaScript 而已。

数组

使用 np.array() 创建一个数组：

import { numpy as np } from "@jax-js/jax";

const ar = np.array([1, 2, 3]);

默认情况下，这是一个 float32 类型的数组，但你可以指定不同的数据类型：

const ar = np.array([1, 2, 3], { dtype: np.int32 });

为了更高效地构建数组，可以从 JavaScript 的 TypedArray 缓冲区创建数组：

const buf = new Float32Array([10, 20, 30, 100, 200, 300]);
const ar = np.array(buf).reshape([2, 3]);

当你完成操作后，可以将 jax.Array 解包回 JavaScript。这也会应用任何待处理的操作或惰性更新：

// 1) 返回可能嵌套的 JavaScript 数组。
ar.js();
await ar.jsAsync(); // 更快，非阻塞

// 2) 返回扁平的 TypedArray 数据缓冲区。
ar.dataSync();
await ar.data(); // 最快，非阻塞

数组可以应用数学运算。例如：

import { numpy as np, scipySpecial as special } from "@jax-js/jax";

const x = np.arange(100).astype(np.float32); // 整数数组 [0..99]

const y1 = x.ref.add(x.ref); // x + x
const y2 = np.sin(x.ref); // sin(x)
const y3 = np.tanh(x.ref).mul(5); // 5 * tanh(x)
const y4 = special.erfc(x.ref); // erfc(x)

请注意，在上述代码中，我们使用了 x.ref。这是因为 jax-js 使用引用计数的“所有权”模型来跟踪数组内存何时可以释放。关于这一点将在下文详细说明。

引用计数

大型数组会占用大量内存。Python 的机器学习库会重写 __del__() 方法来释放内存，但 JavaScript 并没有这样的 API 来运行对象的析构函数（参考）。这意味着你必须手动跟踪引用。jax-js 尝试使这一过程尽可能地便捷，以避免在循环中意外泄漏内存。

每个 jax.Array 都有一个引用计数。它遵循以下规则：

• 每当你创建一个数组时，它的引用计数从 1 开始。
• 当数组的引用计数达到 0 时，它会被释放，且不能再被使用。
• 对于一个数组 a：
  • 访问 a.ref 会返回 a，并将其引用计数加 1。
  • 将 a 作为参数传递给任何函数时，其引用计数会减 1。
  • 调用 a.dispose() 也会使其引用计数减 1。

这意味着 jax-js 中的所有函数都必须以引用的方式“接管”它们的参数。每当你想将一个数组作为参数传递时，可以直接传递它以释放其所有权，或者使用 .ref 如果你希望稍后再使用它。

你也必须在自己的函数中遵守这些规则！ 所有组合函数如 jvp、grad、jit 都假定你按照这些约定传递参数，并且会尊重这些约定。

// 不好：两次使用 `x`，导致其引用计数两次递减。
function foo_bad(x: np.Array, y: np.Array) {
  return x.add(x.mul(y));
}

// 好：第一次使用 `x` 是通过 `x.ref`，引用计数加 `1`。
function foo_good(x: np.Array, y: np.Array) {
  return x.ref.add(x.mul(y));
}

再看一个例子：

// 不好：在 `if` 分支中没有消耗 `x`。
function bar_bad(x: np.Array, skip: boolean) {
  if (skip) return np.zeros(x.shape);
  return x;
}

// 好：在每个分支中只消耗一次 `x`。
function bar_good(x: np.Array, skip: boolean) {
  if (skip) {
    const ret = np.zeros(x.shape);
    x.dispose();
    return ret;
  }
  return x;
}

你可以假设 jax-js 中的每一个函数都会正确地接管所有权，除了少数极罕见且已记录的情况。

grad()、vmap() 和 jit()

JAX 的可组合变换签名也在 jax-js 中得到支持。下面是一个使用 grad 和 vmap 计算函数导数的简单示例：

import { numpy as np, grad, vmap } from "@jax-js/jax";

const x = np.linspace(-10, 10, 1000);

const y1 = vmap(grad(np.sin))(x.ref); // d/dx sin(x) = cos(x)
const y2 = np.cos(x);

np.allclose(y1, y2); // => true

jit 函数在 GPU 上执行长时间的基本操作序列时尤其有用，因为它会将多个操作融合为一次内核调度。这提高了硬件加速器上的内存带宽利用率，而 GPU 上的瓶颈往往不是浮点运算能力，而是内存带宽。例如：

export const hypot = jit(function hypot(x1: np.Array, x2: np.Array) {
  return np.sqrt(np.square(x1).add(np.square(x2)));
});

如果没有 JIT，hypot() 函数需要四次内核调度：两次乘法、一次加法和一次平方根。而 JIT 会将这些操作融合为一个内核，一次性完成所有计算。

所有函数式变换都可以接受输入和输出的类型化 JsTree。它们类似于 JAX 的 pytrees（文档），基本上就是由嵌套的 JavaScript 对象和数组组成的结构。例如：

import { grad, numpy as np } from "@jax-js/jax";

type Params = {
  foo: np.Array;
  bar: np.Array[];
};

function getSums(p: Params) {
  const fooSum = p.foo.sum();
  const barSum = p.bar.map((x) => x.sum()).reduce(np.add);
  return fooSum.add(barSum);
}

grad(getSums)({
  foo: np.array([1, 2, 3]),
  bar: [np.array([10]), np.array([11, 12])],
});
// => { foo: [1, 1, 1], bar: [[1], [1, 1]] }

请注意，你需要使用 type 别名语法而不是 interface 来定义细粒度的 JsTree 类型。

设备

与 JAX 类似，jax-js 也引入了“设备”的概念。设备是后端，用于在内存中存储数组，并决定如何在其上执行编译后的操作。

目前 jax-js 中有 4 种设备：

• cpu：速度较慢的解释型 JavaScript 实现，仅用于调试。
• wasm：WebAssembly，在支持 SharedArrayBuffer 的情况下可实现多线程。
• webgpu：WebGPU，在支持的浏览器（Chrome、Firefox、Safari、iOS）上可用。
• webgl：基于片段着色器的 WebGL2。这是一种较旧的图形 API，几乎可在所有浏览器上运行，但性能远低于 WebGPU。该选项以尽力而为的方式提供，支持程度较低。

我们推荐使用 webgpu 以获得最佳性能，尤其是在运行神经网络时。 默认设备是 wasm，但您可以在启动时更改它：

import { defaultDevice, init } from "@jax-js/jax";

const devices = await init(); // 启动所有可用的后端。

if (devices.includes("webgpu")) {
  defaultDevice("webgpu");
} else {
  console.warn("WebGPU 不受支持，将回退到 Wasm。");
}

您还可以将单个数组放置在特定设备上：

import { devicePut, numpy as np } from "@jax-js/jax";

const ar = np.array([1, 2, 3]); // 默认设备为 "wasm"
await devicePut(ar, "webgpu"); // 现在设备变为 "webgpu"

辅助库

@jax-js 命名空间下还有其他库，既可以与 jax-js 配合使用，也可以在其他项目中独立使用。

• @jax-js/loaders 可从 Safetensors 等多种格式加载张量，包含快速且符合标准的 BPE 实现，并通过 OPFS 缓存大型资产（如模型权重）的 HTTP 请求。
• @jax-js/onnx 是一个将 ONNX 格式模型转换为原生 jax-js 函数的模型加载器。
• @jax-js/optax 提供 Adam 和 SGD 等优化器的实现。

性能

要在浏览器开发者工具中查看每个内核的跟踪信息，请调用 jax.profiler.startTrace()。

WebGPU 运行时包含一个面向 ML 的编译器，具备分块感知优化功能，并针对不同浏览器进行了调优。此外，该库还独有 jit() 功能，可以将多个操作融合在一起并记录执行图。在 Apple M4 Max 芯片上，jax-js 在矩阵乘法运算中可达到 超过 7000 GFLOP/s 的性能（试一试）。

就该示例而言，其速度显著高于使用手写自定义内核库的 TensorFlow.js 和 ONNX Runtime Web。

不过目前仍处于早期阶段，仍有大量易于实现的优化点可供挖掘，同时也有诸如 FlashAttention 变体等独特的优化方向。

API 参考

本次简短教程到此结束。有关详细文档，请参阅生成的API 参考。

开发

以下技术细节适用于参与 jax-js 的贡献以及对其内部进行修改。

本仓库由 pnpm 管理。您可以通过以下命令以监听模式编译和构建所有包：

pnpm install
pnpm run build:watch

随后，您可以使用 Vitest 在无头浏览器中运行测试：

pnpm exec playwright install
pnpm test

目前我们使用的是较旧版本的 Playwright，该版本支持在无头模式下使用 WebGPU；较新版本则会跳过 WebGPU 测试。

要启动运行网站、演示和 REPL 的 Vite 开发服务器：

pnpm -C website dev

您还可以通过以下命令运行代码检查工具、代码格式化工具和类型检查工具：

pnpm lint          # 运行 ESLint
pnpm format        # 使用 Prettier 格式化所有文件
pnpm format:check  # 检查格式而不进行修改
pnpm check         # 运行 TypeScript 类型检查

未来工作 / 招募帮助

欢迎各位贡献！以下是一些值得探索的方向：

• 添加对更多 JAX 函数和操作的支持，详见 兼容性表。
• 提升 WebGPU 和 Wasm 运行时的性能，生成更优的内核，并利用 SIMD 和多线程技术。（即使是单线程的 Wasm，性能也可能提升约 20 倍。）
• 帮助 JIT 编译器在更多场景下融合操作，例如 tanh 分支。
• 构建高效的 Transformer 推理引擎，与 onnxruntime-web 进行对比。

您也可以加入我们的 Discord 社区，与社区成员交流讨论。

jax-js 快速上手指南

jax-js 是一个专为浏览器设计的机器学习框架，旨在将 JAX 风格的高性能 CPU/GPU 计算能力带入 JavaScript 环境。它无需外部依赖，完全在客户端运行，支持 WebAssembly 和 WebGPU 后端。

环境准备

• 系统要求：现代操作系统（Windows, macOS, Linux, iOS, Android）。
• 浏览器支持：
  • CPU (Wasm)：所有主流浏览器（Chrome, Edge, Firefox, Safari）及 Node.js/Deno。
  • GPU (WebGPU)：Chrome/Edge (全平台), Firefox/Safari (macOS 26+/iOS 26+), Chrome for Android。
  • GPU (WebGL)：作为备选方案，支持更广泛的旧版浏览器。
• 前置依赖：无外部依赖库。建议使用支持 ES Modules 的开发环境或现代浏览器。

安装步骤

你可以通过 npm 包管理器安装，或者直接在 HTML 中通过 CDN 引入。

方式一：使用 npm 安装（推荐用于构建项目）

npm i @jax-js/jax

方式二：使用 CDN（无需构建工具）

在 HTML 文件中直接使用 <script type="module"> 引入：

<script type="module">
  import { numpy as np } from "https://esm.sh/@jax-js/jax";
</script>

基本使用

jax-js 的 API 设计与 NumPy/JAX 高度兼容。以下是核心用法示例。

1. 创建与操作数组

导入库并创建数组，执行基本的数学运算：

import { numpy as np } from "@jax-js/jax";

// 创建数组 (默认为 float32)
const x = np.array([1, 2, 3]);

// 执行运算
const y = x.mul(4); // 结果: [4, 8, 12]

// 指定数据类型
const z = np.array([1, 2, 3], { dtype: np.int32 });

// 从 TypedArray 高效创建并重塑形状
const buf = new Float32Array([10, 20, 30, 100, 200, 300]);
const ar = np.array(buf).reshape([2, 3]);

2. 数据导出

将 jax.Array 转换回原生 JavaScript 数据：

// 获取嵌套的 JS 数组
const jsArray = ar.js(); 
// 或者异步非阻塞获取（推荐）
const jsArrayAsync = await ar.jsAsync(); 

// 获取扁平化的 TypedArray 缓冲区
const dataBuffer = ar.dataSync();
// 或者异步非阻塞获取（最快）
const dataBufferAsync = await ar.data(); 

3. 高级特性：自动微分与 JIT 编译

利用 JAX 风格的变换函数进行梯度计算和算子融合：

import { numpy as np, grad, vmap, jit } from "@jax-js/jax";

// 自动微分示例：计算 sin(x) 的导数 (应为 cos(x))
const x = np.linspace(-10, 10, 1000);
const derivative = vmap(grad(np.sin))(x.ref); 
const expected = np.cos(x);
console.log(np.allclose(derivative, expected)); // true

// JIT 编译示例：融合多个算子为单个内核，提升 GPU 性能
const hypot = jit(function hypot(x1, x2) {
  return np.sqrt(np.square(x1).add(np.square(x2)));
});

⚠️ 重要提示：引用计数管理

由于 JavaScript 缺乏确定的析构机制，jax-js 采用引用计数来管理内存。

• 规则：创建数组时引用计数为 1。传入函数参数会消耗引用（-1），访问 .ref 会增加引用（+1）。
• 最佳实践：如果变量需要在当前作用域多次使用，请使用 .ref；如果只使用一次或直接传参，可直接使用变量。

// ✅ 正确用法：第一次使用 x 时调用 .ref 增加计数，第二次使用消耗计数
function safeOp(x, y) {
  return x.ref.add(x.mul(y));
}

// ❌ 错误用法：两次使用 x 导致引用计数提前归零，引发内存错误或崩溃
function unsafeOp(x, y) {
  return x.add(x.mul(y)); 
}