Nested schemas: point-cloud statistics

Note

All examples are expected to run from the examples/<example_name> directory of the Tesseract-Torch repository.

Real Tesseracts rarely take a flat list of arrays. Inputs and outputs are usually grouped into sub-models, some fields are differentiable arrays and others are plain settings such as strings and numbers. Tesseract-Torch’s apply_tesseract() mirrors that structure: you pass nested dictionaries in, you get nested dictionaries out, and autograd flows through whichever leaves are differentiable.

In this example, you will learn how to:

  1. Build a Tesseract whose inputs and outputs are nested Pydantic models.

  2. Call it through apply_tesseract() with nested dictionaries that mix tensors and plain Python values.

  3. Take reverse-mode gradients through a nested output field with .backward().

  4. Do the same in forward mode with torch.autograd.forward_ad.

Step 1: Build + serve example Tesseract

The example Tesseract summarizes a weighted point cloud. Its input schema has two sub-models: cloud holds the differentiable arrays (positions and weights), and settings holds plain options (label and scale). Its output schema is nested as well: statistics carries the differentiable results, summary echoes back non-differentiable bookkeeping. Here is the computation at the heart of it (see pointstats_torch/tesseract_api.py):

def evaluate(inputs: dict) -> dict:
    positions = inputs["settings"]["scale"] * inputs["cloud"]["positions"]
    weights = inputs["cloud"]["weights"]
    total_weight = weights.sum()

    barycenter = (weights[:, None] * positions).sum(dim=0) / total_weight
    offsets = positions - barycenter
    radius_of_gyration = torch.sqrt(
        (weights * (offsets * offsets).sum(dim=1)).sum() / total_weight
    )

    return {
        "statistics": {
            "barycenter": barycenter,
            "radius_of_gyration": radius_of_gyration,
        },
        "summary": {
            "label": inputs["settings"]["label"],
            "n_points": positions.shape[0],
        },
    }

You may build the example Tesseract either via the command line, or running the cell below (you can skip running this if already built).

%%bash
# Build pointstats_torch Tesseract so we can use it below
tesseract build pointstats_torch/

To interact with the Tesseract, we use the Python SDK from tesseract_core to load the built image and start a server container.

from tesseract_core import Tesseract

pointstats = Tesseract.from_image("pointstats_torch")
pointstats.serve()

Step 2: Invoke the Tesseract via Tesseract-Torch

The input dictionary follows the schema one level at a time: cloud and settings are dictionaries themselves, and only the leaves are tensors or plain values. Let’s summarize the four corners of a unit square, all with the same weight:

\[\mathbf{p}_1 = (0, 0, 0), \quad \mathbf{p}_2 = (1, 0, 0), \quad \mathbf{p}_3 = (1, 1, 0), \quad \mathbf{p}_4 = (0, 1, 0)\]

settings takes a string and a float, which apply_tesseract() passes through to the Tesseract untouched.

from pprint import pprint

import torch

from tesseract_torch import apply_tesseract

cloud = {
    "positions": torch.tensor(
        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]]
    ),
    "weights": torch.tensor([1.0, 1.0, 1.0, 1.0]),
}
settings = {"label": "unit square", "scale": 1.0}

outputs = apply_tesseract(pointstats, inputs={"cloud": cloud, "settings": settings})
pprint(outputs)

The output is nested like the schema: outputs["statistics"] holds tensors, outputs["summary"] holds the label and point count as plain Python values. As expected, the barycenter is \((0.5, 0.5, 0)\) and every corner sits at distance \(\sqrt{0.5} \approx 0.707\) from it.

Step 3: Gradients through a nested output

Because the positions are a differentiable leaf, we can take gradients of any function of the nested outputs with respect to them. Autodifferentiation is dispatched to the underlying Tesseract’s vector_jacobian_product endpoint, which only sees the paths that actually need gradients (cloud.positions here, not settings.scale).

# Reverse-mode AD via .backward()
cloud_grad = {
    "positions": torch.tensor(
        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]],
        requires_grad=True,
    ),
    "weights": torch.tensor([1.0, 1.0, 1.0, 1.0]),
}

result = apply_tesseract(pointstats, inputs={"cloud": cloud_grad, "settings": settings})
result["statistics"]["barycenter"].sum().backward()
print("Reverse-mode gradient of cloud['positions']:")
print(cloud_grad["positions"].grad)

Each point contributes its weight divided by the total weight to every coordinate of the barycenter, so with four equal weights the gradient is \(0.25\) everywhere.

Weights are differentiable too. Let’s make them uneven and ask how the radius of gyration responds to each weight, this time via torch.autograd.grad:

cloud_w = {
    "positions": torch.tensor(
        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]]
    ),
    "weights": torch.tensor([4.0, 1.0, 1.0, 1.0], requires_grad=True),
}

result = apply_tesseract(pointstats, inputs={"cloud": cloud_w, "settings": settings})
(grad_w,) = torch.autograd.grad(
    result["statistics"]["radius_of_gyration"], cloud_w["weights"]
)
print("d(radius_of_gyration) / d(weights):")
print(grad_w)

Forward-mode AD works the same way through torch.autograd.forward_ad, dispatching to the Tesseract’s jacobian_vector_product endpoint. Here we push a tangent of ones through the positions and read the tangent off the nested barycenter output:

import torch.autograd.forward_ad as fwAD

tangent = torch.ones_like(cloud["positions"])

with fwAD.dual_level():
    cloud_dual = {
        "positions": fwAD.make_dual(cloud["positions"], tangent),
        "weights": cloud["weights"],
    }
    result = apply_tesseract(
        pointstats, inputs={"cloud": cloud_dual, "settings": settings}
    )
    _, jvp_result = fwAD.unpack_dual(result["statistics"]["barycenter"])

print("Forward-mode JVP of the barycenter:")
print(jvp_result)

Shifting every point by \((1, 1, 1)\) shifts the barycenter by \((1, 1, 1)\), which is exactly what the JVP reports.

Step N+1: Clean-up and conclusions

Since we kept the Tesseract alive using .serve(), we need to manually stop it using .teardown() to avoid leaking resources.

This is not necessary when using Tesseract in a with statement, as it will automatically clean up when the context is exited.

pointstats.teardown()

And that’s it! Nested schemas need no special handling on the PyTorch side: mirror the schema with nested dictionaries, and autograd follows the differentiable leaves wherever they sit in the tree.