jac_to_grad

torchjd.autojac.jac_to_grad(tensors, aggregator, retain_jac=False)

Aggregates the Jacobians stored in the .jac fields of tensors and accumulates the result into their .grad fields.

Parameters:

• tensors (Iterable[Tensor]) – The tensors whose .jac fields should be aggregated. All Jacobians must have the same first dimension (e.g. number of losses).

• aggregator (Aggregator) – The aggregator used to reduce the Jacobians into gradients.

• retain_jac (bool) – Whether to preserve the .jac fields of the tensors after they have been used. Defaults to False.

Return type:

None

Note

This function starts by “flattening” the .jac fields into matrices (i.e. flattening all of their dimensions except the first one), then concatenates those matrices into a combined Jacobian matrix. The aggregator is then used on this matrix, which returns a combined gradient vector, that is split and reshaped to fit into the .grad fields of the tensors.

Example

This example shows how to use jac_to_grad after a call to backward

>>> import torch
>>>
>>> from torchjd.aggregation import UPGrad
>>> from torchjd.autojac import backward, jac_to_grad
>>>
>>> param = torch.tensor([1., 2.], requires_grad=True)
>>> # Compute arbitrary quantities that are function of param
>>> y1 = torch.tensor([-1., 1.]) @ param
>>> y2 = (param ** 2).sum()
>>>
>>> backward([y1, y2])  # param now has a .jac field
>>> jac_to_grad([param], aggregator=UPGrad())  # param now has a .grad field
>>> param.grad
tensor([-1.,  1.])

The .grad field of param now contains the aggregation (by UPGrad) of the Jacobian of \(\begin{bmatrix}y_1 \\ y_2\end{bmatrix}\) with respect to param.