API reference

PyTorch compatibility layer for Tesseract.

Provides apply_tesseract(), which wraps any Tesseract as a differentiable PyTorch operation supporting both reverse-mode (.backward()) and forward-mode (torch.autograd.forward_ad) automatic differentiation.

tesseract_torch.apply_tesseract(tesseract, inputs, *, vmap_method=None)[source]

Call a Tesseract as a differentiable PyTorch operation.

Infers which inputs/outputs are differentiable from the Tesseract’s schema. Torch tensors provided for differentiable fields participate in autograd; all other values are passed through as static inputs.

Supports both reverse-mode (.backward()) and forward-mode (torch.autograd.forward_ad) differentiation.

CUDA tensors stay on the device whenever the Tesseract offers a GPU transport that works from this process, and are copied through the host otherwise. Pass tesseract.with_encoding(gpu_transport="none") to always copy them. GPU transports are an experimental tesseract-core feature. See CUDA tensors.

Parameters:

  • tesseract (Tesseract) – A Tesseract instance.

  • inputs (dict[str, Any]) – Nested dict matching the Tesseract’s input schema. Provide torch.Tensor for array fields you want gradients through, and plain Python / NumPy values for everything else.

  • vmap_method (str | None) – How the call is batched under torch.vmap. None (default) raises if a batched tensor reaches the call. "sequential" calls the Tesseract once per batch element and works with any schema. "expand_dims" gives every unbatched tensor or NumPy array input a leading dimension of size 1 and calls the Tesseract once; its schema must accept the extra leading dimension (e.g. Array[..., Float64]) and it must broadcast size 1 against the batch. "broadcast_all" does the same, but broadcasts the unbatched inputs to the full batch size, for Tesseracts that need matching shapes. Under torch.vmap non-differentiable array outputs are returned as tensors, since a NumPy array cannot carry the batch dimension. A non-array output is returned once: the first element’s value under "sequential", with a warning if another element differs, and the batched call’s value, which describes the whole batch, under the other two methods. See Batching with torch.vmap.

Return type:

dict[str, Any]

Returns:

Nested dict matching the Tesseract’s output schema, with differentiable array outputs as torch.Tensor (with grad_fn when inputs require grad) and non-differentiable outputs as NumPy arrays or scalars. Arrays that arrive over a GPU transport are CUDA tensors instead, and tensors an in-process endpoint returns are passed through. Under torch.vmap non-differentiable array outputs are tensors.

Example:

# Flat schema
result = apply_tesseract(quadratic, {"x": x, "A": A, "b": b})
result["y"].sum().backward()

# Nested schema
result = apply_tesseract(meshstats, {
    "mesh": {"n_points": 3, ..., "points": points_tensor}
})
result["statistics"]["barycenter"].sum().backward()