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

Question

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

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

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

偏导数的顺序无关紧要。

所以，我有几个问题：

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

标签： pythonprobability-densityprobability-distributionautomatic-differentiationjax

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