jax.custom_jvp.defjvp#

custom_jvp.defjvp(jvp, symbolic_zeros=False)[源代码]#

为此实例表示的函数定义自定义 JVP 规则。

参数：

  • jvp (Callable[..., tuple[ReturnValue, ReturnValue]]) – 一个Python可调用对象，表示自定义的JVP规则。当没有nondiff_argnums时，jvp函数应接受两个参数，其中第一个是原始输入的元组，第二个是切线输入的元组。两个元组的长度等于custom_jvp函数的参数数量。jvp函数应生成一个对作为输出，其中第一个元素是原始输出，第二个元素是切线输出。输入和输出元组的元素可以是数组或任何嵌套的元组/列表/字典。

  • symbolic_zeros (bool) – 布尔值，指示是否应在其切线参数中传递表示静态符号零的对象，以对应未扰动的值；否则，仅传递标准的JAX类型（例如，类数组）。将此选项设置为True允许JVP规则检测某些输入是否未参与微分，但代价是需要对这些对象进行特殊处理（例如，不能将它们传递到jax.numpy函数中）。默认值为False。

返回：

返回 jvp，以便可以将defjvp用作装饰器。

返回类型：

Callable[…, tuple[ReturnValue, ReturnValue]]

示例

>>> @jax.custom_jvp
... def f(x, y):
...   return jnp.sin(x) * y
...
>>> @f.defjvp
... def f_jvp(primals, tangents):
...   x, y = primals
...   x_dot, y_dot = tangents
...   primal_out = f(x, y)
...   tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot
...   return primal_out, tangent_out

>>> 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))