JAX对蒙特卡洛看涨期权求交叉二阶导数报非标量输出错误如何解决

• JAX的grad算子仅支持对输出为标量的函数求导，你定义的second_derivative_mc返回的payoff形状为(100,)，不符合grad的输入要求，所以触发类型错误。
• 代码中使用了NumPy的随机数生成接口，NumPy操作不在JAX的计算图跟踪范围内，会导致梯度计算结果错误，甚至无法求导。
• vmap与grad的调用逻辑、函数内部的维度设置也存在不匹配的问题。

1. 调整函数内部维度设置，确保单个标的价格、单个波动率输入时，函数输出为标量期权价格，移除不必要的N=100批量逻辑（批量处理交给vmap完成）
2. 替换NumPy随机数为JAX原生的随机生成接口，保证随机采样过程可被JAX跟踪
3. 调整微分与向量化的调用顺序，先对单输入单输出的函数求交叉二阶导，再用vmap批量处理所有参数对

from jax import jit, grad, vmap
import jax.numpy as jnp
from jax import random
import numpy as np

Underlying_asset = jnp.linspace(1.1,1.4,100)
volatilities = jnp.linspace(0.5,0.6,100)

def second_derivative_mc(S, vol):
    j, T, q, r, k = 10000, 1., 0, 0, 1.
    S0 = S
    # 生成JAX可跟踪的随机数
    key = random.PRNGKey(10)
    U = random.normal(key, (j,)) 
    sigma2 = vol ** 2

    first = sigma2 * jnp.ones(j)
    second = vol * U

    X = -0.5 * first + jnp.sqrt(T) * second

    St = jnp.exp(X) * S0

    P = jnp.maximum(St - k, 0)
    payoff = jnp.average(P) * jnp.exp(-q * T)
    return payoff 

# 先求交叉二阶导，再向量化批量处理
cross_deriv = grad(grad(second_derivative_mc, argnums=1), argnums=0)
greek = vmap(cross_deriv)(Underlying_asset, volatilities)

如果确实需要保留N=100的批量逻辑在函数内部，可以对输出的payoff做降维处理，比如取均值、或者调整argnums的对应关系，保证求导时函数输出为标量即可。

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