企业官网建设流程全解析

JAX 快速上手：Python+NumPy 程序可组合变换完全指南

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/gh_mirrors/jax/jax

JAX 是一个面向 Python 和 NumPy 的数值计算库：你照常用 NumPy 风格写函数，JAX 可组合变换就能在函数外面套上自动求导（grad）、向量化（vmap）和 GPU/TPU 上的 JIT 编译（jit），而且这些变换可以任意互相嵌套组合。本文带你在三分钟内完成安装、跑通第一个 JAX 自动求导示例，并拆清楚 jit 是怎么把 Python 编译到加速器上的。

JAX 能帮你做什么

对没接触过 JAX 的人来说，价值放在具体场景里才看得见：

• 训练模型要梯度？一行jax.grad(loss)自动推导出梯度函数，不用手推链式法则；
• 单个样本的模型想一次跑 1024 个样本？jax.vmap把逐样本函数变成批量函数，不用自己写广播逻辑；
• Python 层循环成了瓶颈、GPU 在空转？jax.jit把整个函数编译后丢给 GPU/TPU，后续调用速度接近原生 C++。

这 trio 组合——grad、vmap、jit——就是 JAX 官方描述里"composable transformations（可组合变换）"的核心。

五分钟跑通第一个 JAX 程序 ⚡

先装，CPU 环境直接：

pip install -U jax

然后打开 Python 交互环境，跑这一段：

import jax import jax.numpy as jnp def mse(params, x, y): return jnp.mean((x @ params - y) ** 2) x = jnp.array([[1.0, 2.0], [3.0, 4.0]]) y = jnp.array([5.0, 11.0]) print(jax.grad(mse)(jnp.zeros(2), x, y))

最后一行直接输出损失对两个参数的梯度，这就"跑通了"。更妙的是：一行代码不改，把jax.jit包在mse外面，同样的计算就编译到 GPU 上执行——JAX JIT 编译对源码零侵入。

核心能力拆解

grad：一行代码拿到自动求导梯度

jax.grad不是靠多次运行函数去估算梯度，而是把函数"追踪"（trace）一遍，推导出精确的梯度函数。它支持高阶导数、Hessian、Jacobian，遇到jax.lax.cond、while_loop这类带控制流的函数也能照常求导，基本不需要你手动处理分支情况。

vmap：把逐样本函数变批量函数

jax.vmap(f)把"一个一个算"变成"一整批并行算"，批量部分由硬件的向量化指令承担，比 Python for 循环通常快一个数量级。它还支持嵌套结构（pytree）输入和多批次维度，写模型时经常拿来自动 batch。

jit：把 Python 一键编译到 GPU/TPU

jax.jit会先把函数 trace 成中间表示 Jaxpr，再交给 XLA 编译器生成目标硬件可执行的代码，顺带融合小算子、省掉冗余的内存读写。

"追踪 → 变换 → 编译"整条链路对用户完全透明，这也是 JAX 写起来"无感"的来源：变换叠加的顺序不同，得到的编译产物也不同，但你的函数体永远只有一份。

进阶：多设备与性能调优场景怎么玩

多设备并行。模型大到单卡放不下时，用jax.pmap或jax.shard_map把程序切分并复制到多台设备；XLA 的 SPMD 机制保证整段程序只编译一次、在 N 台设备上以同一份代码并行执行。

性能剖析。觉得哪里慢了，开启jax.profiler生成时间线，用 Perfetto UI 打开，哪个 Python 函数、哪个 XLA 算子在吃时间一眼可见，调优有的放矢。

多机多卡的完整玩法和调优细节，建议深挖 docs/notebooks/Distributed_arrays_and_automatic_parallelization.md 和 docs/profiling.md。

写在最后

一句话总结：JAX 让你继续写熟悉的 Python，把求导、批量、编译、并行的脏活全交给变换和编译器。把 examples/ 里的样例今天跑一遍，梯度、vmap、jit 的手感一下午就能建立起来 ✅

• 官方文档：docs/
• 示例代码：examples/
• TPU 云上手 notebook：cloud_tpu_colabs/

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/gh_mirrors/jax/jax