JAX 中无需通过 custom_jvp 声明非可微参数；正确做法是在调用 jax.grad 等自动微分变换时，显式指定 argnums 参数，仅对需微分的变量求导，其余参数天然视为常量（即非可微）。这既简洁又符合 JAX 的设计哲学。
jax 中无需通过 `custom_jvp` 声明非可微参数；正确做法是在调用 `jax.grad` 等自动微分变换时，显式指定 `argnums` 参数，仅对需微分的变量求导，其余参数天然视为常量（即非可微）。这既简洁又符合 jax 的设计哲学。
在 JAX 中，一个常见误解是：必须借助 @jax.custom_jvp 并设置 nondiff_argnums 来“标记”某些参数为非可微。但事实上，nondiff_argnums 并非用于声明“哪些参数不参与求导”，而是用于在自定义前向微分（JVP）逻辑中，明确哪些参数不参与线性化计算——它本质是梯度定制机制的一部分，而非通用的参数过滤开关。
真正的、通用且推荐的做法是：将非可微参数作为普通 Python 参数传入函数，而在调用 jax.grad、jax.jacobian、jax.vjp 等高阶变换时，通过 argnums（或 has_aux、holomorphic 等）精确控制对哪些位置的参数求导。未被选中的参数在自动微分过程中自动被视为常量（constant），其值不会被追踪、不会生成梯度，也不会触发不可微错误（除非它们意外进入可微子图）。
例如：
import jax
import jax.numpy as jnp
def linear_layer(x, W, b, activation="relu"):
# activation 是字符串，显然不可微；W 和 b 是数组，可微
y = x @ W + b
if activation == "relu":
return jnp.maximum(0.0, y)
elif activation == "sigmoid":
return jax.nn.sigmoid(y)
else:
raise ValueError("Unsupported activation")
# ✅ 正确：仅对第 1、2 个参数（W 和 b）求梯度，activation（第 3 个）被忽略
grad_fn = jax.grad(linear_layer, argnums=(1, 2))
x = jnp.ones((3, 4))
W = jnp.ones((4, 2)) * 0.1
b = jnp.zeros(2)
loss = lambda W, b: jnp.sum(linear_layer(x, W, b, activation="relu"))
gW, gb = grad_fn(x, W, b, "relu") # 注意：x 和 activation 传入但不求导
⚠️ 注意事项：
-
argnums接受整数或整数元组，索引基于函数实际调用时的参数顺序（而非nondiff_argnums所指的“定义时”顺序）； - 若函数含
*args或**kwargs，应避免依赖位置索引，建议重构为显式命名参数并使用argnums明确指定； -
custom_jvp仅在你需要重写某段逻辑的梯度行为（如实现稳定数值梯度、绕过不可微操作、注入近似梯度）时才需使用；为“省事不写梯度”而强行套用custom_jvp反而违背 JAX 的函数式与显式原则，且易引发Tracer错误或静默失败； - 对于
jit、vmap等变换，argnums同样适用，且与grad组合时具有完全正交性。
总结：JAX 的自动微分系统是调用侧驱动（call-site driven） 的——梯度责任由 grad() 的调用者定义，而非函数定义者强制约束。因此，最通用、最安全、最符合 JAX 惯例的方式，永远是：在 jax.grad(..., argnums=...) 中精准声明求导目标，其余参数自然“绝缘”于微分图之外。
