tensorplay.func.jacrev

tensorplay.func.jacrev(func: Callable, argnums: int | tuple[int, ...] = 0, *, has_aux: bool = False, chunk_size: int | None = None, _preallocate_and_copy: bool = False)

Returns a function computing the Jacobian of func by reverse mode.

Reverse mode costs one backward pass per output element, so jacrev is the cheaper direction when the output is smaller than the input.

Parameters:

• func (Callable) – a Python function returning one or more tensors.

• argnums (int or Tuple[int]) – which positional arguments to differentiate with respect to. Default: 0.

• has_aux (bool) – whether func returns (output, aux).

• chunk_size (int, optional) – compute the Jacobian chunk_size rows at a time to bound peak memory. None computes it in one shot.

Example

>>> jacobian = jacrev(tensorplay.sin)(tensorplay.randn(5))
>>> jacobian.shape
tensorplay.Size(5, 5)