jax.custom_vjp.defvjp

jax.custom_vjp.defvjp#

custom_vjp.defvjp(fwd, bwd, symbolic_zeros=False, optimize_remat=False)[source]#

为该实例表示的函数定义一个自定义 VJP 规则。

参数：:

  • fwd (Callable[..., tuple[ReturnValue, Any]]) – 代表自定义 VJP 规则的前向传递的 Python 可调用对象。当没有 nondiff_argnums 时，fwd 函数具有与底层原始函数相同的输入签名。它应返回一个输出对，其中第一个元素表示原始输出，第二个元素表示要从前向传递中存储的任何“残差”值，以便在后向传递中由函数 bwd 使用。输入参数和输出对的元素可以是数组或其嵌套元组/列表/字典。

  • bwd (Callable[..., tuple[Any, ...]]) – 表示自定义 VJP 规则的反向传播的 Python 可调用对象。当没有 nondiff_argnums 时，bwd 函数接受两个参数，第一个参数是 fwd 在正向传播中产生的“残差”值，第二个参数是与原始函数输出结构相同的输出余切。 bwd 的输出必须是长度等于原始函数参数数量的元组，并且元组元素可以是数组或嵌套的元组/列表/字典，以便与原始输入参数的结构匹配。

  • symbolic_zeros (bool) –

    布尔值，用于确定是否向 fwd 和 bwd 规则指示符号零。启用此选项允许自定义导数规则检测何时某些输入以及何时某些输出余切不参与微分。如果 True

    • fwd 必须接受，代替构成原始函数参数的 pytree 中每个叶值 x，一个具有两个属性的对象（类型为 jax.custom_derivatives.CustomVJPPrimal）：value 和 perturbed。 value 字段是原始主参数，而 perturbed 是一个布尔值。 perturbed 位指示参数是否参与微分（即，如果它为 False，则相应的雅可比矩阵“列”为零）。

    • bwd 将在其余切参数中传递表示静态符号零的对象，以对应于未扰动值；否则，只传递标准 JAX 类型（例如，数组）。

    将此选项设置为 True 允许这些规则检测某些输入和输出是否不参与微分，但代价是需要特殊处理。例如

    • fwd 的签名会发生变化，它传递的对象不能直接从规则中输出。

    • bwd 规则传递的对象并非完全是数组，也不能传递给大多数 jax.numpy 函数。

    • 参与主函数参数的任何自定义 pytree 节点必须在其展开函数中接受作为 fwd 规则输入叶传递的两个字段记录对象。

    默认值为 False。

  • optimize_remat (bool) – 布尔值，一个实验性标志，用于在 jax.remat() 下使用此函数时启用自动优化。当 fwd 规则是诸如 Pallas 内核或自定义调用之类的隐藏调用时，这将非常有用。默认值为 False。

返回值：:

无。

返回类型：:

无

示例

>>> @jax.custom_vjp
... def f(x, y):
...   return jnp.sin(x) * y
...
>>> def f_fwd(x, y):
...   return f(x, y), (jnp.cos(x), jnp.sin(x), y)
...
>>> def f_bwd(res, g):
...   cos_x, sin_x, y = res
...   return (cos_x * g * y, sin_x * g)
...
>>> f.defvjp(f_fwd, f_bwd)

>>> x = jnp.float32(1.0)
>>> y = jnp.float32(2.0)
>>> with jnp.printoptions(precision=2):
...   print(jax.value_and_grad(f)(x, y))
(Array(1.68, dtype=float32), Array(1.08, dtype=float32))
目录