如何从 Jax 中的联合累积密度函数计算联合概率密度函数？

Question

我在 python 中定义了一个联合累积密度函数作为 jax 数组的函数并返回单个值。 比如：

def cumulative(inputs: array) -> float:
    ...

要获得梯度，我知道我可以做grad(cumulative)，但这只是给我累积相对于输入变量的一阶偏导数。 相反，我想做的是计算这个，假设 F 是我的函数，f 是联合概率密度函数：

偏导的顺序无关紧要。

所以，我有几个问题：

• 如何在 Jax 中高效计算？我想我不能只打电话给 grad n 次
• 一旦计算出结果函数，结果函数的调用复杂度是否会比原始函数更高（是增加了 O(n)，还是常数，还是其他什么）？
• 或者，我如何计算仅关于输入数组的一个变量而不是整个数组的单个偏导数？ （我将重复此 n 次，每个变量一次）

Answer 1

JAX 通常将渐变视为相对于单个参数，而不是参数中的元素。在这种情况下，一个与您想要做的类似（但不完全相同）的内置函数是 jax.hessian，它计算二阶导数的 hessian 矩阵；例如：

import jax
import jax.numpy as jnp

def f(x):
  return jnp.prod(x ** 2)

x = jnp.arange(1.0, 4.0)
print(jax.hessian(f)(x))
# [[72. 72. 48.]
#  [72. 18. 24.]
#  [48. 24.  8.]]

对于数组中单个元素的高阶导数，我认为您必须手动嵌套渐变。您可以使用如下所示的辅助函数来执行此操作：

def grad_all(f):
  def gradfun(x):
    args = tuple(x)
    f_args = lambda *args: f(jnp.array(args))
    for i in range(len(args)):
      f_args = jax.grad(f_args, argnums=i)
    return f_args(*args)
  return gradfun

print(grad_all(f)(x))
# 48.0