jacobian_trace#

bayesflow.utils.jacobian_trace(f: Callable[[Tensor], Tensor], x: Tensor, max_steps: int = None, return_output: bool = False, seed: int | SeedGenerator = None)[source]#

Compute or estimate the trace of the Jacobian matrix of f.

Parameters:
fcallable

The function to be differentiated.

xTensor of shape (n, …, d)

The input tensor to f.

max_stepsint, optional

The maximum number of steps to use for the estimate. If this does not exceed the dimensionality of f(x), use Hutchinson’s algorithm to return an unbiased estimate of the Jacobian trace. Otherwise, perform an exact computation. Default: None

return_outputbool, optional

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

seedint or keras SeedGenerator, optional

The seed to use for hutchinson trace estimation. Only has an effect when max_steps < d.

Returns:
2-tuple of tensors:

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

  2. Tensor of shape (n,)

    An unbiased estimate or the exact trace of the Jacobian of f.