本文深入探讨了在JAX中为自定义类（如PyTorch风格的Module）计算梯度时遇到的常见问题及其解决方案。核心挑战在于jax.grad要求被微分函数以参数形式接收可微分变量，并需要JAX识别这些复杂对象内部的参数结构。文章将介绍如何重构损失函数以适应jax.grad，并强调JAX PyTree机制的重要性，包括手动注册和利用Flax/Equinox等框架进行参数管理的最佳实践，从而实现对自定义模型权重的有效梯度计算。
JAX梯度计算的核心挑战
在使用jax进行机器学习模型开发时，我们经常会构建类似pytorch nn.module的自定义类来封装模型结构和参数。然而，当尝试直接使用jax.grad来计算这些自定义类中权重（例如线性层的weights和biases）的梯度时，可能会发现jax.grad只返回一个与损失相关的单一梯度，而非模型内部参数的梯度。这通常是因为jax.grad在处理复杂对象时，需要满足两个基本条件：
- 可微分参数的函数化: jax.grad作用于一个函数，该函数必须将其需要微分的参数作为显式输入。
- 参数结构的JAX识别: JAX需要理解复杂对象（如自定义Module实例）的内部结构，以识别哪些部分是可微分的参数。
原始代码示例中，criterion函数直接接收模型的输出，而不是模型本身及其参数。此外，JAX默认不了解自定义Model类内部的linear.weights和linear.biases是需要微分的参数。
解决方案一：重构损失函数以适应jax.grad
jax.grad的设计理念是作用于一个纯函数，并计算该函数关于其特定参数的梯度。因此，我们需要将模型的调用和损失计算封装在一个新的函数中，该函数以模型实例（或其参数）作为输入。
考虑以下原始模型和损失函数定义：
import jax import jax.numpy as jnp from jax.tree_util import register_pytree_node # 提前导入，用于后续PyTree注册 class Module: def __init__(self) -> None: pass def __call__(self, inputs: jax.