jacobian#

bayesflow.utils.jacobian(f: Callable[[Tensor], Tensor], x: Tensor, return_output: bool = False)[source]#

Compute the Jacobian matrix of f with respect to x.

Parameters:
fcallable

The function to be differentiated.

xTensor of shape (…, D_in)

The input tensor to f.

return_outputbool, optional

Whether to return the output of f(x) along with the Jacobian matrix. Default: False

Returns:
Tensor of shape (…, D_out, D_in)

The Jacobian matrix of f with respect to x.

2-tuple of tensors

  1. The output of f(x) (if return_output is True)

  2. Tensor of shape (…, D_out, D_in)

    The Jacobian matrix of f with respect to x.