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. 
