使用JAX的algorithmic differentiation求函数对a、b偏导报错求助

核心错误原因 #

• 符号微分实现逻辑错误：原代码中dfb的计算依然指定对a求导，未正确设置求导目标为b
• JAX实现混用不兼容依赖：JAX自动微分要求待求导函数使用jax.numpy提供的数学算子，不能调用sympy的sym.cos/sym.exp等接口
• 函数定义不符合JAX规范：待求导函数foo未声明输入参数a、b，无法接收求导输入值
• 求导调用逻辑错误：
  • jax.grad默认仅对函数的第一个输入参数求导，要同时获取对a、b的偏导需指定argnums=(0,1)
  • 原代码直接返回jax.grad(foo)生成的函数对象，没有传入具体的a、b值计算偏导结果，触发报错：

cannot unpack non-iterable function object

修正后实现代码 #

符号微分版本 #

import sympy as sym

# 预先定义符号变量与目标函数表达式
a_sym, b_sym = sym.symbols('a b')
f_expr = (2 / b_sym) * sym.cos(a_sym) * sym.exp(- (a_sym / b_sym) ** 2)

def get_partial_derivatives_sym(a_val: float, b_val: float) -> tuple[float, float]:
    # 分别对a、b求偏导后代入数值计算
    dfa = sym.diff(f_expr, a_sym).evalf(subs={a_sym: a_val, b_sym: b_val})
    dfb = sym.diff(f_expr, b_sym).evalf(subs={a_sym: a_val, b_sym: b_val})
    return (float(dfa), float(dfb))

JAX算法微分版本 #

import jax
import jax.numpy as jnp

# 定义待求导的原函数，全部使用jax.numpy提供的数学算子
def f(a: jnp.ndarray, b: jnp.ndarray) -> jnp.ndarray:
    x = (2 / b) * jnp.cos(a)
    y = jnp.exp(- jnp.power(a / b, 2))
    return x * y

# 生成同时对第0位参数(a)、第1位参数(b)求偏导的函数
f_grad = jax.grad(f, argnums=(0, 1))

def get_partial_derivatives_jax(a_val: float, b_val: float) -> tuple[float, float]:
    # 传入具体数值计算偏导，返回要求的元组格式
    dfa, dfb = f_grad(a_val, b_val)
    return (float(dfa), float(dfb))

调用示例 #

# 测试输入a=1.0，b=2.0
print(get_partial_derivatives_jax(1.0, 2.0))

内容的提问来源于stack exchange，提问作者Aaron