python - JAX - 区分功能的问题

Question

我正在尝试在调用中执行蒙特卡罗模拟，然后在 Python 中计算其相对于基础资产的一阶导数，但它仍然不起作用

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

xi = jnp.linspace(1,1.2,5)
def Simulation(xi):
    K,T,number_sim,sigma,r,q = 1.,1.,100,0.4,0,0
    S = jnp.broadcast_to(xi,(number_sim,len(xi))).T

    mean = -.5 * sigma * sigma * T
    volatility = sigma*jnp.sqrt(T)
    r_numb = random.PRNGKey(10)
    BM = mean + volatility * random.normal(r_numb, shape=(number_sim,))

    product = S*jnp.exp(BM)

    payoff = jnp.maximum(product-K,0)

    result = jnp.average(payoff, axis=1)*jnp.exp(-q*T)

    return result

first_derivative = vmap(grad(Simulation))(xi)

我不知道实现该算法的方式是否是使用“AD 方法”计算导数的最佳方式；该算法以这种方式工作：

• S = 模拟一个包含所有底层证券的矩阵；对于每一行，我使用“xi = jnp.linspace”生成每个底层，并且在矩阵的每一行内，我有相同的值多次等于“number_sim”

• product = 生成 BM（包含正常数的向量）后，我需要将BM的每个元素（带有 exp）与S的每一行的每个元素相乘

所以这是对算法的简短解释，我非常感谢任何管理这个问题的建议或技巧，并用 AD 方法计算导数！提前致谢

Answer 1