jax.vjp

Contents

jax.vjp#

jax.vjp(fun: Callable[..., T], *primals: Any, has_aux: Literal[False] = False, reduce_axes: Sequence[AxisName] = (), saveable_args: Any = True, in_nzs: Any = None) → tuple[T, Callable][source]#
jax.vjp(fun: Callable[..., tuple[T, U]], *primals: Any, has_aux: Literal[True], reduce_axes: Sequence[AxisName] = (), saveable_args: Any = True, in_nzs: Any = None) → tuple[T, Callable, U]

Compute a (reverse-mode) vector-Jacobian product of fun.

grad() is implemented as a special case of vjp().

Parameters:

  • fun – Function to be differentiated. Its arguments should be arrays, scalars, or standard Python containers of arrays or scalars. It should return an array, scalar, or standard Python container of arrays or scalars.

  • primals – A sequence of primal values at which the Jacobian of fun should be evaluated. The number of primals should be equal to the number of positional parameters of fun. Each primal value should be an array, a scalar, or a pytree (standard Python containers) thereof.

  • has_aux – Optional, bool. Indicates whether fun returns a pair where the first element is considered the output of the mathematical function to be differentiated and the second element is auxiliary data. Default False.

  • saveable_args – Optional, a tuple-tree of bools (i.e. nested tuples with bool leaves) or equivalently a pytree prefix of the primals with bool leaves, by default the single bool True. Indicates whether each primal argument (or argument sub-pytree, or leaf) may be saved for the backward pass. It must form a tree prefix of primals up to pytree node types: tuples are matched against argument containers only by their number of children, so e.g. a tuple entry can correspond to a dict argument. Where a False entry applies, argument values that would have been saved verbatim as residuals are instead replaced by NotSaveable sentinels in the args_res attribute of vjpfun, and the caller must restore them (e.g. by assigning to vjpfun.args_res) before applying vjpfun. Only argument values saved verbatim are affected; residuals computed from the arguments are saved as usual.

  • in_nzs – Optional, a tuple-tree of bools like saveable_args, by default None meaning all-True. Declares which primal inputs have (possibly) nonzero tangents. Where a False entry applies, that input’s tangent is treated as symbolically zero during linearization, which can make more outputs’ tangents symbolically zero; the resulting per-output nonzeros pattern is available as the out_nzs attribute of vjpfun, and vjpfun returns a zero cotangent for any False-marked input.

Returns:

If has_aux is False, returns a (primals_out, vjpfun) pair, where primals_out is fun(*primals). If has_aux is True, returns a (primals_out, vjpfun, aux) tuple where aux is the auxiliary data returned by fun.

vjpfun is a function from a cotangent vector with the same shape as primals_out to a tuple of cotangent vectors with the same number and shapes as primals, representing the vector-Jacobian product of fun evaluated at primals.

>>> import jax
>>>
>>> def f(x, y):
...   return jax.numpy.sin(x), jax.numpy.cos(y)
...
>>> primals, f_vjp = jax.vjp(f, 0.5, 1.0)
>>> xbar, ybar = f_vjp((-0.7, 0.3))
>>> print(xbar)
-0.61430776
>>> print(ybar)
-0.2524413
Contents