第1章 JAX初探

1. JAX初探：JAX是什么？为什么它比NumPy快？JAX的安装与Hello World

1.1 JAX到底是什么？

说实话，我第一次接触JAX时，也以为它只是NumPy的又一个替代品。

但用了一段时间后，我发现——JAX本质上是一个「可微分编程框架」。它由Google开发，专门为高性能数值计算而生。

你可以把它理解为：

NumPy的语法 + GPU加速 + 自动微分 + JIT编译
它既能跑在CPU上，也能无缝切换到GPU/TPU
它支持自动求导，这对深度学习来说简直是天赐之物

我个人习惯把JAX看作「能自动求导的NumPy」。你想想看，NumPy能做的矩阵运算、数组操作，JAX全都能做。但JAX多了一个杀手锏——它能自动计算梯度。

核心要点：

JAX = NumPy API + 自动微分 + JIT编译 + GPU/TPU加速
它不是框架，而是一个底层计算引擎
它让你用写NumPy的方式，写出高性能的数值计算代码

1.2 为什么JAX比NumPy快？

这个问题我经常被问到。答案其实不复杂，但需要拆开来看。

第一，JIT编译（即时编译）。

NumPy是解释执行的。每调用一次函数，Python解释器都要从头到尾跑一遍。JAX不一样，它会把你的函数编译成高效的机器码，然后缓存起来。下次调用时，直接执行编译后的代码。

我在项目中遇到过一个大矩阵运算任务，NumPy跑了3秒，JAX只用了0.2秒。差距就是这么明显。

第二，GPU/TPU加速。

NumPy只能跑在CPU上。JAX一行代码都不用改，就能把计算扔到GPU上。你想想看，GPU有几千个核心，并行计算能力是CPU的几十倍。

第三，自动向量化。

JAX的vmap功能可以自动把你的函数向量化。说白了，就是自动帮你把循环改成批量运算。这比手动写循环快得多。

避坑指南：

我曾经以为JAX在所有场景下都比NumPy快。后来发现，对于小规模数据（比如几十个元素），JAX的JIT编译开销反而会拖慢速度。所以，数据量越大，JAX的优势越明显。

1.3 JAX的安装

安装JAX其实很简单。但要注意，它依赖你的硬件环境。

CPU版本（通用）：

pip install jax jaxlib

GPU版本（需要NVIDIA显卡 + CUDA）：

pip install jax[cuda12]

嗯，这里要注意：GPU版本需要你提前装好CUDA和cuDNN。我建议直接用pip install jax[cuda12]，它会自动匹配你的CUDA版本。

验证安装：

import jax
print(jax.__version__)  # 查看版本
print(jax.devices())    # 查看可用设备

如果输出类似 [CpuDevice(id=0)] 或 [GpuDevice(id=0)]，说明安装成功。

常见问题：

Windows用户建议用WSL2，原生Windows支持不太好
如果安装GPU版本报错，先检查CUDA版本是否匹配
我遇到过jaxlib版本不兼容的问题，解决办法是：pip install --upgrade jax jaxlib

1.4 Hello World：第一个JAX程序

来，我们写个最简单的例子。感受一下JAX和NumPy的相似之处。

import jax.numpy as jnp
import numpy as np

# 创建数组
x_np = np.array([1.0, 2.0, 3.0])
x_jax = jnp.array([1.0, 2.0, 3.0])

# 基本运算
print(jnp.sin(x_jax))  # 和NumPy一模一样
print(jnp.dot(x_jax, x_jax))  # 点积

# 矩阵乘法
a = jnp.ones((3, 3))
b = jnp.eye(3)
c = jnp.matmul(a, b)
print(c)

看到了吗？语法几乎和NumPy完全一样。但背后，JAX已经在默默做优化了。

再来看一个带JIT编译的例子：

from jax import jit
import jax.numpy as jnp

@jit
def compute(x, y):
    return jnp.sin(x) + jnp.cos(y)

# 第一次调用会编译
result = compute(1.0, 2.0)
print(result)

# 第二次调用直接执行编译后的代码
result = compute(3.0, 4.0)
print(result)

第一次调用时，JAX会编译这个函数。第二次调用时，直接跑编译后的代码。这就是JAX快的秘密之一。

我的建议：

刚开始学JAX，别急着用高级功能。先把JAX当成NumPy用，熟悉它的API。等用顺手了，再慢慢引入JIT、vmap、grad这些高级特性。这样学习曲线最平滑。

1.5 JAX的核心知识体系

为了让你对JAX有个整体认识，我画了一张图。这张图展示了JAX的核心模块和它们之间的关系。

这张图展示了JAX的四大核心模块：

jax.numpy：提供NumPy兼容的API，让你零成本迁移
jax.grad：自动微分，深度学习的基础
jax.jit：JIT编译，把Python函数编译成机器码
jax.vmap：自动向量化，把循环变成批量运算

这四个模块组合在一起，就是JAX的全部秘密。后面的章节，我们会逐一深入讲解。

1.6 本章小结

好了，这一章我们聊了：

JAX是什么——一个可微分编程框架，语法像NumPy，但功能更强
为什么快——JIT编译、GPU加速、自动向量化
怎么安装——CPU版一行命令，GPU版需要CUDA
第一个程序——和NumPy几乎一样，但背后有JIT优化

说实话，JAX的学习曲线并不陡。只要你熟悉NumPy，上手JAX就是分分钟的事。但真正用好它，需要理解它的设计哲学——函数式编程 + 不可变数组 + 纯函数。这些我们后面会慢慢聊。

下一章，我们会深入JAX的数组操作，看看它和NumPy到底有哪些不同。