{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": "# Basic example: vector addition\n\n
\n

Note

\n\nAll examples are expected to run from the `examples/` directory of the [Tesseract-Torch repository](https://github.com/pasteurlabs/tesseract-torch).\n
\n\nTesseract-Torch is a lightweight extension to [Tesseract Core](https://github.com/pasteurlabs/tesseract-core) that wraps Tesseracts as differentiable [PyTorch](https://pytorch.org/) operations, with full support for reverse-mode and forward-mode automatic differentiation.\n\nIn this example, you will learn how to:\n1. Build a Tesseract that performs vector addition.\n1. Access its endpoints via Tesseract-Torch's `apply_tesseract()` function.\n1. Compose Tesseracts into more complex functions, blending multiple Tesseract applications with local PyTorch operations.\n2. Autodifferentiate the resulting pipeline via `.backward()` (reverse-mode) and `torch.autograd.forward_ad` (forward-mode)." }, { "cell_type": "markdown", "metadata": {}, "source": "## Step 1: Build + serve example Tesseract\n\nIn 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`):\n\n```python\ndef evaluate(inputs: dict) -> dict:\n a_scaled = inputs[\"a\"][\"s\"] * inputs[\"a\"][\"v\"]\n b_scaled = inputs[\"b\"][\"s\"] * inputs[\"b\"][\"v\"]\n add_result = a_scaled + b_scaled\n min_result = a_scaled - b_scaled\n\n def safe_norm(x, ord):\n return torch.pow(torch.pow(torch.abs(x), ord).sum() + 1e-8, 1.0 / ord)\n\n return {\n \"vector_add\": {\n \"result\": add_result,\n \"normed_result\": add_result / safe_norm(add_result, ord=inputs[\"norm_ord\"]),\n },\n \"vector_min\": {\n \"result\": min_result,\n \"normed_result\": min_result / safe_norm(min_result, ord=inputs[\"norm_ord\"]),\n },\n }\n```\n\nYou may build the example Tesseract either via the command line, or running the cell below (you can skip running this if already built)." }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "%%bash\n# Build vectoradd_torch Tesseract so we can use it below\ntesseract build vectoradd_torch/" }, { "cell_type": "markdown", "metadata": {}, "source": [ "To interact with the Tesseract, we use the Python SDK from `tesseract_core` to load the built image and start a server container." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "from tesseract_core import Tesseract\n\nvectoradd = Tesseract.from_image(\"vectoradd_torch\")\nvectoradd.serve()" }, { "cell_type": "markdown", "metadata": {}, "source": "## Step 2: Invoke the Tesseract via Tesseract-Torch\n\nUsing the `vectoradd_torch` Tesseract image we built earlier, let's add two vectors together, representing the following operation:\n\n$$\\begin{pmatrix} 1 \\\\ 2 \\\\ 3 \\end{pmatrix} + 2 \\cdot \\begin{pmatrix} 4 \\\\ 5 \\\\ 6 \\end{pmatrix} = \\begin{pmatrix} 9 \\\\ 12 \\\\ 15 \\end{pmatrix}$$\n\nWe 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." }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "from pprint import pprint\n\nimport torch\n\nfrom tesseract_torch import apply_tesseract\n\na = {\"v\": torch.tensor([1.0, 2.0, 3.0])}\nb = {\n \"v\": torch.tensor([4.0, 5.0, 6.0]),\n \"s\": torch.tensor(2.0),\n}\n\noutputs = apply_tesseract(vectoradd, inputs={\"a\": a, \"b\": b})\npprint(outputs)" }, { "cell_type": "markdown", "metadata": {}, "source": [ "As expected, `outputs['vector_add']` gives a value of $(9, 12, 15)$." ] }, { "cell_type": "markdown", "metadata": {}, "source": "## Step 3: Function composition via Tesseracts\n\nTesseract-Torch enables you to compose chains of Tesseract evaluations, blended with local PyTorch operations, while retaining full autograd support.\n\nThe 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." }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def fancy_operation(a, b):\n \"\"\"Fancy operation.\"\"\"\n result = apply_tesseract(vectoradd, inputs={\"a\": a, \"b\": b})\n result = apply_tesseract(\n vectoradd, inputs={\"a\": {\"v\": result[\"vector_add\"][\"result\"]}, \"b\": b}\n )\n # We can mix and match with local PyTorch operations\n result = 2.0 * result[\"vector_add\"][\"normed_result\"] + b[\"v\"]\n result = apply_tesseract(vectoradd, inputs={\"a\": {\"v\": result}, \"b\": b})\n return result[\"vector_add\"][\"result\"][1]\n\n\nfancy_operation(a, b)" }, { "cell_type": "markdown", "metadata": {}, "source": "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`):" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Reverse-mode AD via .backward()\na_grad = {\"v\": torch.tensor([1.0, 2.0, 3.0], requires_grad=True)}\nb_nograd = {\n \"v\": torch.tensor([4.0, 5.0, 6.0]),\n \"s\": torch.tensor(2.0),\n}\n\nloss = fancy_operation(a_grad, b_nograd)\nloss.backward()\nprint(\"Reverse-mode gradient of a['v']:\")\nprint(a_grad[\"v\"].grad)" }, { "cell_type": "markdown", "metadata": {}, "source": "We can also use `torch.autograd.grad` for more control, and forward-mode AD via `torch.autograd.forward_ad`:" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import torch.autograd.forward_ad as fwAD\n\n# torch.autograd.grad for reverse-mode autodiff\na_grad = {\"v\": torch.tensor([1.0, 2.0, 3.0], requires_grad=True)}\nloss = fancy_operation(a_grad, b_nograd)\n(grad_a,) = torch.autograd.grad(loss, a_grad[\"v\"])\nprint(\"torch.autograd.grad result:\")\nprint(grad_a)\n\n# Forward-mode AD via torch.autograd.forward_ad\na_fwd = {\"v\": torch.tensor([1.0, 2.0, 3.0])}\ntangent = torch.ones_like(a_fwd[\"v\"])\n\nwith fwAD.dual_level():\n a_fwd_dual = {\"v\": fwAD.make_dual(a_fwd[\"v\"], tangent)}\n result = apply_tesseract(vectoradd, inputs={\"a\": a_fwd_dual, \"b\": b_nograd})\n _, jvp_result = fwAD.unpack_dual(result[\"vector_add\"][\"result\"])\n\nprint(\"\\nForward-mode JVP result:\")\nprint(jvp_result)" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step N+1: Clean-up and conclusions" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Since we kept the Tesseract alive using `.serve()`, we need to manually stop it using `.teardown()` to avoid leaking resources. \n", "\n", "This is not necessary when using `Tesseract` in a `with` statement, as it will automatically clean up when the context is exited." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "vectoradd.teardown()" }, { "cell_type": "markdown", "metadata": {}, "source": "And that's it!\nYou've learned how to build up differentiable pipelines with Tesseracts that integrate seamlessly with PyTorch's autograd." } ], "metadata": { "kernelspec": { "display_name": "science", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.7" } }, "nbformat": 4, "nbformat_minor": 4 }