saiunit.autograd.jacobian

Contents

saiunit.autograd.jacobian#

saiunit.autograd.jacobian(fun, argnums=0, has_aux=False, holomorphic=False, allow_int=False)[source]#

Alias of jacrev().

This is a convenience alias that delegates directly to jacrev(). See jacrev() for full documentation.

Parameters:

  • fun (Callable) – Function whose Jacobian is to be computed.

  • argnums (int | Sequence[int]) – Specifies which positional argument(s) to differentiate with respect to. Default is 0.

  • has_aux (bool) – If True, fun returns (output, aux) and only output is differentiated. Default is False.

  • holomorphic (bool) – Whether fun is promised to be holomorphic. Default is False.

  • allow_int (bool) – Whether integer-valued inputs are allowed. Default is False.

Returns:

jacfun – A function that computes the Jacobian of fun through reverse-mode automatic differentiation.

Return type:

Callable

See also

jacrev

The primary implementation.

jacfwd

Forward-mode Jacobian computation.

Examples

>>> import jax.numpy as jnp
>>> import saiunit as u
>>> import saiunit.autograd as suauto
>>> def f(x):
...     return x ** 2
>>> jac_fn = suauto.jacobian(f)
>>> jac_fn(jnp.array(3.0) * u.ms)
6.0 * ms
Contents