tfep.utils.math.batch_autograd_jacobian

tfep.utils.math.batch_autograd_jacobian(x, y, **grad_kwargs)

Compute the batch jacobian of y w.r.t. x using torch.autograd.grad.

Parameters:

• x (torch.Tensor) – Shape (batch_size, *). Input.

• y (torch.Tensor) – Shape (batch_size, *). Output.

• **grad_kwargs – Keyword arguments for autograd.grad (e.g., create_graph, retain_graph).

Returns:

jac – Shape (batch_size, size_y, size_x).

Return type:

torch.Tensor