2

我在我的项目中经常使用 VJP。它运行受雅可比计算约束的函数，并返回一个 primals_out 以及可调用的 vjp 函数。例如，JAX 文档中的自定义 VJP 定义如下所示：

from jax import custom_vjp

@custom_vjp
def f(x, y):
  return jnp.sin(x) * y

def f_fwd(x, y):
# Returns primal output and residuals to be used in backward pass by f_bwd.
  return f(x, y), (jnp.cos(x), jnp.sin(x), y)

def f_bwd(res, g):
  cos_x, sin_x, y = res # Gets residuals computed in f_fwd
  return (cos_x * g * y, sin_x * g)

f.defvjp(f_fwd, f_bwd)

在此示例中，我们看到使用 VJP 时需要评估前向函数。使用常规 VJP 而不是自定义 VJP 时也是如此。但是，当函数的评估成本很高并且因为我已经在我的代码中的某个地方运行了该函数时，我不希望 VJP 再评估一次该函数。

那么，有没有办法表明在计算其 VJP 时不会评估函数？

4

1 回答 1

1

我认为在这种情况下没有任何方法可以显式禁用前向评估，但是如果您将计算包装在jit编译中，XLA 编译器将自动执行死代码消除并从计算图中修剪未使用的分支。

于 2021-07-23T14:25:57.157 回答