JAX基础操作：从数组到随机数，一个工程师的实战笔记
大家好，我是你们的老朋友。今天我们来聊聊JAX的基础操作。
说实话，我第一次接触JAX时，心里想的是：「又一个深度学习框架？我学不动了。」但真正用起来才发现，JAX的设计哲学非常优雅——它把NumPy的简洁和自动微分的强大结合在了一起。嗯，这很对我胃口。
核心要点：JAX的数组操作和NumPy几乎一模一样，但背后多了自动微分和GPU加速的能力。你只需要改变一行导入代码，就能获得飞一般的速度。
1. JAX数组创建与基本运算
先说说数组创建。我个人习惯用 jax.numpy 来替代 numpy，因为这样写出来的代码可以直接在GPU上跑。
import jax
import jax.numpy as jnp
# 创建数组
a = jnp.array([1, 2, 3, 4, 5])
b = jnp.zeros((3, 4))
c = jnp.ones((2, 3))
d = jnp.arange(10)
e = jnp.linspace(0, 1, 5)
print(a) # [1 2 3 4 5]
print(b) # 3x4的全零矩阵
基本运算也很直观。加减乘除、矩阵乘法，写法跟NumPy一模一样。
x = jnp.array([[1, 2], [3, 4]])
y = jnp.array([[5, 6], [7, 8]])
# 逐元素运算
print(x + y) # 加法
print(x * y) # 逐元素乘法
# 矩阵乘法
print(jnp.dot(x, y)) # 或者 x @ y
小技巧：我在项目中遇到过一个问题——用 jnp.dot 做矩阵乘法时，如果矩阵维度不匹配，报错信息比NumPy更详细。这点很贴心。
2. JAX与NumPy的互操作性
你想想看，如果JAX和NumPy不能互通，那得多麻烦？好在JAX设计者早就想到了这一点。
JAX数组可以轻松转成NumPy数组，反过来也行。但有个坑——JAX数组默认在GPU上，而NumPy数组在CPU上。转换时会有数据拷贝的开销。
import numpy as np
# JAX -> NumPy
jax_array = jnp.array([1, 2, 3])
numpy_array = np.array(jax_array) # 或者 jax_array.__array__()
# NumPy -> JAX
numpy_data = np.random.randn(3, 4)
jax_data = jnp.array(numpy_data)
避坑指南：我曾经在训练循环里频繁做JAX和NumPy的转换，结果性能直接掉了一个数量级。记住：尽量在JAX的世界里完成所有操作，只在最后输出结果时转成NumPy。
还有一个实用函数——jax.device_put()。你可以显式地把数据放到指定设备上。
# 把数据放到GPU上
gpu_data = jax.device_put(numpy_data, jax.devices('gpu')[0])
3. 随机数生成与种子管理
随机数生成这块，JAX和NumPy有本质区别。NumPy用全局随机状态，而JAX要求你显式管理随机种子。为什么？因为JAX的函数式编程风格要求「无副作用」——每次调用随机函数，你都得传入一个状态，然后返回新的状态和随机数。
说白了，就是JAX不允许有隐藏的全局状态。这刚开始让我很不习惯，但用久了发现，这样反而更容易复现实验结果。
from jax import random
# 创建随机种子
key = random.PRNGKey(42)
# 生成随机数
rand_array = random.normal(key, shape=(3, 4))
print(rand_array)
# 生成新的种子
key, subkey = random.split(key)
rand_array2 = random.uniform(subkey, shape=(2, 3))
重要：永远不要重复使用同一个key！每次生成随机数后，都要用 random.split() 创建新key。否则你会得到完全相同的随机数序列——这在训练模型时会导致灾难性的结果。
我建议你养成这样的习惯：
# 好的做法
key = random.PRNGKey(0)
for epoch in range(10):
key, subkey = random.split(key)
batch = random.randint(subkey, shape=(32,), minval=0, maxval=100)
# 训练代码...
为什么会这样设计？因为JAX的 jit 编译和 vmap 向量化要求函数是纯函数——同样的输入必须产生同样的输出。如果随机数生成依赖全局状态，那编译后的函数就没法保证确定性了。
知识体系总览
下面这张图帮你理清本章的核心逻辑：
我的经验：刚开始用JAX时，我总是不自觉地用NumPy的随机函数，结果在jit编译时报错。后来我写了个小脚本，把所有 np.random 替换成 jax.random，并加上key管理逻辑。这个脚本我现在还在用。
好了，这一章的内容就到这里。记住：JAX的数组操作和NumPy几乎一样，但随机数管理完全不同。把key管理好，你的代码就能在GPU上飞起来。