Jacrev

Contents

Jacrev#

nabla.transforms.jacrev.jacrev(func, argnums=0, has_aux=False, holomorphic=False, allow_int=False)[source]#

Compute the Jacobian of a function using reverse-mode autodiff.

Parameters:

  • func (Callable[[...], Any]) – Function to differentiate (should take positional arguments)

  • argnums (int | tuple[int, ...] | list[int]) – Optional, integer or sequence of integers. Specifies which positional argument(s) to differentiate with respect to (default 0).

  • has_aux (bool) – Optional, bool. Indicates whether func 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.

  • holomorphic (bool) – Optional, bool. Indicates whether func is promised to be holomorphic. Default False. Currently ignored.

  • allow_int (bool) – Optional, bool. Whether to allow differentiating with respect to integer valued inputs. Currently ignored.

Returns:

A function with the same arguments as func, that evaluates the Jacobian of func using reverse-mode automatic differentiation. If has_aux is True then a pair of (jacobian, auxiliary_data) is returned.

Return type:

Callable[[…], Any]

Note

This follows JAX’s jacrev API: - Only accepts positional arguments - For functions requiring keyword arguments, use functools.partial or lambda - Returns the Jacobian as a pytree structure matching the input structure

Contents