JAX 简介 #

┌─────────────────────────────────────────────────────────────┐
│                         JAX                                  │
├─────────────────────────────────────────────────────────────┤
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐         │
│  │  自动微分    │  │  自动向量化  │  │  JIT 编译   │         │
│  │   (grad)    │  │   (vmap)    │  │   (jit)     │         │
│  └─────────────┘  └─────────────┘  └─────────────┘         │
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐         │
│  │  NumPy API  │  │  函数式设计  │  │  多设备支持  │         │
│  └─────────────┘  └─────────────┘  └─────────────┘         │
└─────────────────────────────────────────────────────────────┘

2018年 ─── JAX 项目启动
    │
    │      Google Research 内部项目
    │      结合 Autograd 和 XLA
    │
2019年 ─── 开源发布
    │
    │      GitHub 开源
    │      社区开始关注
    │
2020年 ─── 生态系统发展
    │
    │      Flax、Haiku 等库发布
    │      研究论文广泛使用
    │
2021年 ─── 功能增强
    │
    │      pjit 分布式支持
    │      TPU 支持完善
    │
至今   ─── 广泛应用
    │
    │      DeepMind 大量使用
    │      学术研究首选框架
    │      大模型训练

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sum(x ** 2)

x = jnp.array([1.0, 2.0, 3.0])

grad_f = jax.grad(f)           f_grad = grad_f(x)             
vmap_f = jax.vmap(f)           batch_result = vmap_f(batch_x) 
jit_f = jax.jit(f)             jit_result = jit_f(x)          

jit_grad_f = jax.jit(jax.grad(f))

import jax.numpy as jnp

x = jnp.array([1, 2, 3])
y = jnp.zeros((3, 4))
z = jnp.dot(x, y)

def loss(params, x, y):
    predict = jnp.dot(x, params)
    return jnp.mean((predict - y) ** 2)

grad_loss = jax.grad(loss)
gradients = grad_loss(params, x, y)

def process_single(x):
    return jnp.sum(x ** 2)

process_batch = jax.vmap(process_single)
batch_x = jnp.array([[1, 2], [3, 4], [5, 6]])
results = process_batch(batch_x)  

@jax.jit
def fast_function(x):
    return jnp.dot(x, x.T)

result = fast_function(x)  

┌─────────────────────────────────────────────────────────────┐
│                    传统框架的问题                            │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  1. 性能优化复杂                                             │
│     - 需要手动优化                                           │
│     - 编译和执行分离                                         │
│                                                             │
│  2. 批处理繁琐                                               │
│     - 手动编写批处理代码                                     │
│     - 代码重复                                               │
│                                                             │
│  3. 高阶微分困难                                             │
│     - 二阶导数实现复杂                                       │
│     - 性能不佳                                               │
│                                                             │
│  4. 设备迁移成本高                                           │
│     - CPU/GPU/TPU 代码不同                                   │
│     - 需要重写                                              │
│                                                             │
└─────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────┐
│                    JAX 的解决方案                            │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  ✅ JIT 自动优化                                             │
│     - 一行代码启用编译                                       │
│     - XLA 自动优化                                          │
│                                                             │
│  ✅ vmap 自动向量化                                          │
│     - 无需修改原函数                                         │
│     - 自动批处理                                            │
│                                                             │
│  ✅ 高阶微分简单                                             │
│     - grad 可以嵌套                                         │
│     - grad(grad(f)) 即可                                    │
│                                                             │
│  ✅ 设备无关代码                                             │
│     - 同一代码多设备运行                                     │
│     - 无需修改                                              │
│                                                             │
└─────────────────────────────────────────────────────────────┘

@jax.jit
def neural_ode(params, x, t):
    """神经 ODE 研究"""
    def dynamics(x, t):
        return jnp.dot(params, x)
    return jax.experimental.ode.odeint(dynamics, x, t)

@jax.pmap
def train_step(params, batch):
    """多设备并行训练"""
    grads = jax.grad(loss)(params, batch)
    return jax.tree_map(lambda p, g: p - 0.01 * g, params, grads)

def simulate(params, state):
    """物理模拟"""
    return jax.lax.fori_loop(0, 1000, step_fn, state)

def diffusion_step(params, x, t):
    """扩散模型"""
    return jax.grad(score_function)(params, x, t)

✅ 可组合的函数变换
   - grad、vmap、jit 可以任意组合
   - 代码简洁优雅

✅ 高性能
   - JIT 编译优化
   - XLA 后端加速

✅ 多设备支持
   - CPU、GPU、TPU 统一 API
   - 分布式训练简单

✅ 函数式设计
   - 纯函数易于推理
   - 易于测试和调试

✅ 高阶微分
   - 轻松计算高阶导数
   - 适合研究场景

⚠️ 学习曲线
   - 函数式编程思维
   - 需要理解变换概念

⚠️ 生态系统
   - 比 PyTorch 小
   - 部分领域库不完善

⚠️ 生产部署
   - 不如 TensorFlow 成熟
   - 服务化支持有限

⚠️ 副作用处理
   - 需要理解纯函数
   - 状态管理需要技巧

入门阶段
├── JAX 简介（本文）
├── 安装与配置
├── JAX 基础概念
└── NumPy 兼容性

基础阶段
├── 自动微分 (grad)
├── 自动向量化 (vmap)
├── JIT 编译
└── 函数变换组合

进阶阶段
├── 数组操作
├── 线性代数
├── 随机数生成
└── 控制流

高级阶段
├── 构建神经网络
├── 状态管理
├── 训练循环
└── 模型保存与加载