4、函数式编程基础：纯函数的概念、副作用与JAX的设计哲学

说实话，我第一次接触JAX的时候，心里是有点懵的。

明明写的是Python，怎么感觉像在写一种全新的语言？

后来我才明白，JAX骨子里流的不是Python的血，而是函数式编程的血。它要求你写的每个函数都得是「纯函数」——这个概念，就是今天我们要啃的硬骨头。

4.1 纯函数：到底「纯」在哪？

纯函数这个概念，说白了就两条规则：

• 相同的输入，永远返回相同的输出
• 不修改函数外部的任何状态

听起来简单吧？但实际写代码的时候，我们天天都在违反这两条。

举个例子，看看下面这个函数：

# 不纯的函数
global_counter = 0

def impure_add(x):
    global global_counter
    global_counter += 1
    return x + global_counter

print(impure_add(5))  # 输出 6
print(impure_add(5))  # 输出 7

看到了吗？同样的输入 5，两次调用结果不一样。这就是不纯——它依赖了外部变量 global_counter。

纯函数版本应该是这样的：

# 纯函数
def pure_add(x, counter):
    return x + counter, counter + 1

print(pure_add(5, 0))  # 输出 (6, 1)
print(pure_add(5, 0))  # 输出 (6, 1)

每次调用 pure_add(5, 0)，结果都是 (6, 1)。这才叫纯。

4.2 副作用：代码里的「隐形炸弹」

副作用这个词，听起来挺吓人。其实它指的就是：函数除了返回值之外，对外部世界造成的任何影响。

常见的副作用包括：

• 修改全局变量
• 打印日志到控制台
• 写入文件或数据库
• 修改传入的可变对象（比如列表、字典）
• 发起网络请求

我早期做深度学习项目时，就踩过副作用的坑。有一次训练一个模型，每次跑的结果都不一样。查了两天，最后发现是某个函数里偷偷修改了一个全局的随机种子。嗯，那两天我基本没睡好觉。

看看这个例子：

# 有副作用的函数
def add_to_list(lst, item):
    lst.append(item)  # 修改了传入的列表
    return lst

my_list = [1, 2, 3]
add_to_list(my_list, 4)
print(my_list)  # 输出 [1, 2, 3, 4] —— 原始列表被改了！

这个函数虽然返回了结果，但它偷偷修改了传入的 my_list。这就是副作用。

纯函数版本：

# 无副作用的函数
def add_to_list_pure(lst, item):
    return lst + [item]  # 创建新列表，不修改原列表

my_list = [1, 2, 3]
new_list = add_to_list_pure(my_list, 4)
print(my_list)   # 输出 [1, 2, 3] —— 原列表完好无损
print(new_list)  # 输出 [1, 2, 3, 4]

4.3 JAX的设计哲学：为什么非要「无状态」？

你可能会问：JAX为什么这么较真？非要我们写纯函数？

原因其实很直接——为了自动微分和并行计算。

JAX的核心功能是 grad()、vmap()、pmap() 这些变换。它们的工作原理是：分析你的函数，然后把它转换成另一种形式。

如果函数有副作用，JAX就没法安全地做这种转换。你想想看，如果函数里偷偷修改了全局变量，JAX在求梯度时怎么知道这个变量该不该被微分？

我画了一张图，帮你理解JAX的设计哲学：

说白了，JAX的整个体系都建立在「函数是纯的」这个假设上。一旦你打破了它，所有的高级功能都会出问题。

4.4 如何在JAX中实践纯函数

好了，理论说完了。咱们来点实际的——在JAX里怎么写纯函数。

规则一：不要用全局变量

import jax.numpy as jnp
import jax

# ❌ 错误写法
learning_rate = 0.01

def bad_update(params, grads):
    return params - learning_rate * grads  # 依赖全局变量

# ✅ 正确写法
def good_update(params, grads, learning_rate):
    return params - learning_rate * grads  # 所有依赖都通过参数传入

规则二：不要修改传入的数组

# ❌ 错误写法
def bad_scale(arr, factor):
    arr *= factor  # 修改了传入的数组
    return arr

# ✅ 正确写法
def good_scale(arr, factor):
    return arr * factor  # 返回新数组

规则三：使用JAX的纯函数工具

JAX提供了一些工具，帮你写出更纯的代码：

工具	作用	示例
jax.random.PRNGKey	显式传递随机状态	key = jax.random.PRNGKey(42)
jax.lax.scan	纯函数版本的循环	替代Python的for循环
jax.tree_util	处理嵌套数据结构	不修改原始树结构

4.5 一个完整的例子：纯函数版的线性回归

最后，咱们看一个完整的例子。感受一下纯函数风格和普通风格的区别。

import jax
import jax.numpy as jnp

# 纯函数风格的线性回归
def init_params(key):
    """初始化参数——纯函数"""
    w = jax.random.normal(key, (1,))
    b = jnp.zeros(1)
    return {'w': w, 'b': b}

def predict(params, x):
    """预测——纯函数"""
    return params['w'] * x + params['b']

def loss_fn(params, x, y):
    """损失函数——纯函数"""
    pred = predict(params, x)
    return jnp.mean((pred - y) ** 2)

def update(params, x, y, lr):
    """更新参数——纯函数"""
    grads = jax.grad(loss_fn)(params, x, y)
    return {
        'w': params['w'] - lr * grads['w'],
        'b': params['b'] - lr * grads['b']
    }

# 使用
key = jax.random.PRNGKey(0)
params = init_params(key)
x = jnp.array([1.0, 2.0, 3.0])
y = jnp.array([2.0, 4.0, 6.0])

for _ in range(100):
    params = update(params, x, y, 0.01)

print(params)  # 应该接近 w=2.0, b=0.0

注意看，每个函数都没有副作用。它们只依赖传入的参数，只返回新的值。这就是JAX喜欢的方式。

---

公众号:蓝海资料掘金营，微信deep3321