jax.custom_jvp#
- class jax.custom_jvp(fun, nondiff_argnums=())[source]#
为自定义 JVP 规则定义设置 JAX 可变换函数。
此类旨在用作函数装饰器。实例是可调用对象,其行为类似于应用装饰器的底层函数,但当应用微分变换(例如
jax.jvp()
或jax.grad()
)时除外,在这种情况下,将使用自定义用户提供的 JVP 规则函数,而不是跟踪和执行底层函数实现的自动微分。有两个实例方法可用于定义自定义 JVP 规则:
defjvp()
用于为所有函数输入定义一个单个自定义 JVP 规则,以及为了方便而使用的defjvps()
,它包装了defjvp()
,并允许您为函数相对于其每个参数的偏导数提供单独的定义。例如
@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
有关更详细的介绍,请参阅教程。
- 参数:
fun (Callable[..., ReturnValue])
nondiff_argnums (Sequence[int])
- __init__(fun, nondiff_argnums=())[source]#
- 参数:
fun (Callable[..., ReturnValue])
nondiff_argnums (Sequence[int])
方法
__init__
(fun[, nondiff_argnums])defjvp
(jvp[, symbolic_zeros])为该实例表示的函数定义自定义 JVP 规则。
defjvps
(*jvps)用于分别为每个参数定义 JVP 的便捷包装器。
属性
jvp
symbolic_zeros
fun
nondiff_argnums