Basic example: vector addition¶

Note

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

Tesseract-Torch is a lightweight extension to Tesseract Core that wraps Tesseracts as differentiable PyTorch operations, with full support for reverse-mode and forward-mode automatic differentiation.

In this example, you will learn how to:

Build a Tesseract that performs vector addition.
Access its endpoints via Tesseract-Torch’s apply_tesseract() function.
Compose Tesseracts into more complex functions, blending multiple Tesseract applications with local PyTorch operations.
Autodifferentiate the resulting pipeline via .backward() (reverse-mode) and torch.autograd.forward_ad (forward-mode).

Step 1: Build + serve example Tesseract¶

In this example, we build and use a Tesseract that performs vector addition. The example Tesseract takes two vectors and scalars as input and return some statistics as output. Here is the functionality that’s implemented in the Tesseract (see vectoradd_torch/tesseract_api.py):

def evaluate(inputs: dict) -> dict:
    a_scaled = inputs["a"]["s"] * inputs["a"]["v"]
    b_scaled = inputs["b"]["s"] * inputs["b"]["v"]
    add_result = a_scaled + b_scaled
    min_result = a_scaled - b_scaled

    def safe_norm(x, ord):
        return torch.pow(torch.pow(torch.abs(x), ord).sum() + 1e-8, 1.0 / ord)

    return {
        "vector_add": {
            "result": add_result,
            "normed_result": add_result / safe_norm(add_result, ord=inputs["norm_ord"]),
        },
        "vector_min": {
            "result": min_result,
            "normed_result": min_result / safe_norm(min_result, ord=inputs["norm_ord"]),
        },
    }

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 vectoradd_torch Tesseract so we can use it below
tesseract build vectoradd_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

vectoradd = Tesseract.from_image("vectoradd_torch")
vectoradd.serve()

Step 2: Invoke the Tesseract via Tesseract-Torch¶

Using the vectoradd_torch Tesseract image we built earlier, let’s add two vectors together, representing the following operation:

\[\begin{split}\begin{pmatrix} 1 \\ 2 \\ 3 \end{pmatrix} + 2 \cdot \begin{pmatrix} 4 \\ 5 \\ 6 \end{pmatrix} = \begin{pmatrix} 9 \\ 12 \\ 15 \end{pmatrix}\end{split}\]

We can perform this calculation using the function tesseract_torch.apply_tesseract(), by passing the Tesseract object and the input data as a nested dictionary of PyTorch tensors as inputs.

from pprint import pprint

import torch

from tesseract_torch import apply_tesseract

a = {"v": torch.tensor([1.0, 2.0, 3.0])}
b = {
    "v": torch.tensor([4.0, 5.0, 6.0]),
    "s": torch.tensor(2.0),
}

outputs = apply_tesseract(vectoradd, inputs={"a": a, "b": b})
pprint(outputs)

As expected, outputs['vector_add'] gives a value of \((9, 12, 15)\).

Step 3: Function composition via Tesseracts¶

Tesseract-Torch enables you to compose chains of Tesseract evaluations, blended with local PyTorch operations, while retaining full autograd support.

The function below applies vectoradd twice, ie. \((\mathbf{a} + \mathbf{b}) + \mathbf{a}\), then performs local arithmetic on the outputs, applies vectoradd once more, and finally returns a single element of the result. The resulting function is fully auto-differentiable via PyTorch’s autograd.

def fancy_operation(a, b):
    """Fancy operation."""
    result = apply_tesseract(vectoradd, inputs={"a": a, "b": b})
    result = apply_tesseract(
        vectoradd, inputs={"a": {"v": result["vector_add"]["result"]}, "b": b}
    )
    # We can mix and match with local PyTorch operations
    result = 2.0 * result["vector_add"]["normed_result"] + b["v"]
    result = apply_tesseract(vectoradd, inputs={"a": {"v": result}, "b": b})
    return result["vector_add"]["result"][1]


fancy_operation(a, b)

Autodifferentiation is automatically dispatched to the underlying Tesseract’s jacobian_vector_product and vector_jacobian_product endpoints. Let’s use reverse-mode AD (.backward()) and forward-mode AD (torch.autograd.forward_ad):

# Reverse-mode AD via .backward()
a_grad = {"v": torch.tensor([1.0, 2.0, 3.0], requires_grad=True)}
b_nograd = {
    "v": torch.tensor([4.0, 5.0, 6.0]),
    "s": torch.tensor(2.0),
}

loss = fancy_operation(a_grad, b_nograd)
loss.backward()
print("Reverse-mode gradient of a['v']:")
print(a_grad["v"].grad)

We can also use torch.autograd.grad for more control, and forward-mode AD via torch.autograd.forward_ad:

import torch.autograd.forward_ad as fwAD

# torch.autograd.grad for reverse-mode autodiff
a_grad = {"v": torch.tensor([1.0, 2.0, 3.0], requires_grad=True)}
loss = fancy_operation(a_grad, b_nograd)
(grad_a,) = torch.autograd.grad(loss, a_grad["v"])
print("torch.autograd.grad result:")
print(grad_a)

# Forward-mode AD via torch.autograd.forward_ad
a_fwd = {"v": torch.tensor([1.0, 2.0, 3.0])}
tangent = torch.ones_like(a_fwd["v"])

with fwAD.dual_level():
    a_fwd_dual = {"v": fwAD.make_dual(a_fwd["v"], tangent)}
    result = apply_tesseract(vectoradd, inputs={"a": a_fwd_dual, "b": b_nograd})
    _, jvp_result = fwAD.unpack_dual(result["vector_add"]["result"])

print("\nForward-mode JVP result:")
print(jvp_result)

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.

vectoradd.teardown()

And that’s it! You’ve learned how to build up differentiable pipelines with Tesseracts that integrate seamlessly with PyTorch’s autograd.