jax.custom_jvp.defjvps

custom_jvp.defjvps(*jvps)

用於為每個引數分別定義 JVP 的便利包裝器。

此便利包裝器不能與 nondiff_argnums 一起使用。

參數:

*jvps (Callable[..., ReturnValue] | None) – 一系列函式，每個函式對應 custom_jvp 函式的一個位置引數。每個函式將對應的原始輸入的切線值、原始輸出和ß原始輸入作為引數。請參閱以下範例。

傳回:

無。

傳回類型:

無

範例

>>> @jax.custom_jvp
... def f(x, y):
...   return jnp.sin(x) * y
...
>>> f.defjvps(lambda x_dot, primal_out, x, y: jnp.cos(x) * x_dot * y,
...           lambda y_dot, primal_out, x, y: jnp.sin(x) * y_dot)

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