warp.autograd.jacobian#

warp.autograd.jacobian(
function,
dim=None,
inputs=None,
outputs=None,
input_output_mask=None,
device=None,
max_blocks=0,
block_dim=256,
max_outputs_per_var=-1,
plot_jacobians=False,
metadata=None,
)[source]#

Computes the Jacobians of a function or Warp kernel for the provided selection of differentiable inputs to differentiable outputs.

The input function can be either a Warp kernel (e.g. a function decorated by @wp.kernel) or a regular Python function that accepts arguments (of which some must be Warp arrays) and returns a Warp array or a list of Warp arrays.

In case function is a Warp kernel, its adjoint kernel is launched with the given inputs and outputs, as well as the provided dim, max_blocks, and block_dim arguments (see warp.launch() for more details).

Note

If function is a Warp kernel, the input arguments must precede the output arguments in the kernel code definition.

Only Warp arrays with requires_grad=True are considered for the Jacobian computation.

Function arguments of type wp.struct are not yet supported.

Parameters:

  • function (Kernel | Callable) – The Warp kernel function, or a regular Python function that returns a Warp array or a list of Warp arrays.

  • dim (tuple[int] | None) – The number of threads to launch the kernel, can be an integer, or a Tuple of ints. Only required if function is a Warp kernel.

  • inputs (Sequence | None) – List of input variables. At least one of the arguments must be a Warp array with requires_grad=True.

  • outputs (Sequence | None) – List of output variables. Optional if the function is a regular Python function that returns a Warp array or a list of Warp arrays. Only required if function is a Warp kernel.

  • input_output_mask (list[tuple[str | int, str | int]] | None) – List of tuples specifying the input-output pairs to compute the Jacobian for. Inputs and outputs can be identified either by their integer indices of where they appear in the kernel input/output arguments, or by the respective argument names as strings. If None, computes the Jacobian for all input-output pairs.

  • device (Device | str | None) – The device to launch on (optional). Only used if function is a Warp kernel.

  • max_blocks – The maximum number of CUDA thread blocks to use. Only used if function is a Warp kernel.

  • block_dim – The number of threads per block. Only used if function is a Warp kernel.

  • max_outputs_per_var – Maximum number of output dimensions over which to evaluate the Jacobians for the input-output pairs. Evaluates all output dimensions if value <= 0.

  • plot_jacobians – If True, visualizes the computed Jacobians in a plot (requires matplotlib).

  • metadata (FunctionMetadata | None) – The metadata of the kernel function, containing the input and output labels, strides, and dtypes. If None or empty, the metadata is inferred from the kernel or function.

Returns:

A dictionary of Jacobians, where the keys are tuples of input and output indices, and the values are the Jacobian matrices.

Return type:

dict[tuple[int, int], array]