Autograd Utility: Loop-based Jacobian

Autograd Utility: Loop-based Jacobian#

These derivative functions do not use functorch, but construct the Jacobian row-by-row. This is slower than using functorch.

tad_mctc.autograd.nonfunctorch.jac(a, b, create_graph=None, retain_graph=True)[source]#

Compute the Jacobian of a with respect to b row-by-row.

Parameters:
  • a (Tensor) – Variable that is differentiated.

  • b (Tensor) – Variable with respect to which the derivative is taken.

  • create_graph (bool | None, optional) – Whether to create a backprogatable graph. Required for additional (higher) derivatives. Defaults to True.

  • retain_graph (bool, optional) – Whether to use the multiple graph multiple times. Defaults to True. Otherwise, the graph is deleted after the first call.

Returns:

Jacobian of a with respect to b.

Return type:

Tensor