Hi all,
We’ve just released Tesseract-Torch, a lightweight extension to Tesseract Core that wraps any Tesseract as a differentiable PyTorch operation. If you’ve been using Tesseract-JAX, this is the PyTorch counterpart, and it follows the same design. That means Tesseract is now natively compatible with both JAX and PyTorch, and you can slot a Tesseract into either framework’s autograd graph.
What it does
The whole API is a single function, apply_tesseract, which slots a served Tesseract into PyTorch’s autograd graph:
import torch
from tesseract_core import Tesseract
from tesseract_torch import apply_tesseract
t = Tesseract.from_image("vectoradd_torch")
t.serve()
x = torch.ones(1000, requires_grad=True)
result = apply_tesseract(t, {"a": {"v": x}, "b": {"v": torch.ones(1000)}})
result["vector_add"]["result"].sum().backward()
print(x.grad) # gradients flow back through the Tesseract's VJP endpoint
Both reverse-mode (.backward(), torch.autograd.grad) and forward-mode (torch.autograd.forward_ad) AD are supported. The Tesseract runs in its own container; the gradients come back over HTTP via its vector_jacobian_product / jacobian_vector_product endpoints. From the caller’s side it behaves like any other autograd layer.
Why it might matter to you
Until now, integrating a Tesseract into an end-to-end differentiable program meant working in JAX. PyTorch is where a lot of model code already lives, so this opens the same workflow to that ecosystem: you can keep your network in native PyTorch and call a containerized solver as one layer in the middle of a training loop, with gradients flowing through both.
That’s useful when the solver isn’t something you can just import, for example a heavyweight CFD code, a legacy in-house simulator, or something with a hand-written or Enzyme-generated adjoint. You wrap it once as a Tesseract, then treat it as a differentiable building block.
A worked example
To make this concrete, there’s a new learned closure demo that trains a neural viscosity model end-to-end through a containerized 1D Burgers’ equation solver. The solver lives in a Docker image and exposes a VJP; the closure network is plain torch.nn. Training differentiates through the entire time-stepping loop, with each step making an HTTP call out to the solver container, and recovers a spatially-varying viscosity profile from solution data alone. It’s a small example, but it shows the pattern that the JAX closure demos use, now in PyTorch.
Worth knowing up front
This is a 0.1.0 release with a deliberately small surface, so expect a few sharp edges.
Notably, torch.func transforms (vjp, jvp, grad, vmap) are not supported, because their functionalized tensors can’t be converted to NumPy. (You’ll get a loud error if you try.)
Getting started
pip install tesseract-torch
Feedback is very welcome, especially on the API and on which sharp edges bite you in practice. If you try it on a real problem, we’d love to hear how it goes.