Skip to content

Repository files navigation

Tesseract-Torch

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.

Read the docs | Explore the examples | Report an issue | Talk to the community | Contribute


The API of Tesseract-Torch consists of a single function, apply_tesseract(tesseract, inputs), which integrates any Tesseract into PyTorch's autograd graph:

result = apply_tesseract(my_tesseract, {"x": x_tensor})
result["y"].sum().backward()  # reverse-mode AD
x_tensor.grad  # gradients flow through the Tesseract

Quick start

Note

Before proceeding, make sure you have a working installation of Docker and a modern Python installation (Python 3.10+).

Important

For more detailed installation instructions, please refer to the Tesseract Core documentation.

  1. Install Tesseract-Torch:

    $ pip install tesseract-torch
  2. Build an example Tesseract:

    $ git clone https://github.com/pasteurlabs/tesseract-torch
    $ tesseract build tesseract-torch/examples/simple/vectoradd_torch
  3. Use it as part of a PyTorch program via apply_tesseract:

    import torch
    from tesseract_core import Tesseract
    from tesseract_torch import apply_tesseract
    
    # Load the Tesseract
    t = Tesseract.from_image("vectoradd_torch")
    t.serve()
    
    # Run it with PyTorch tensors
    x = torch.ones(1000, requires_grad=True)
    y = torch.ones(1000)
    
    def vector_sum(x, y):
        res = apply_tesseract(t, {"a": {"v": x}, "b": {"v": y}})
        return res["vector_add"]["result"].sum()
    
    loss = vector_sum(x, y)
    loss.backward()
    print(x.grad)  # gradients via the Tesseract's VJP endpoint
    
    # Forward-mode AD is also supported via torch.autograd.forward_ad

Tip

Now you're ready to jump into our examples for more ways to use Tesseract-Torch.

Sharp edges

  • Required endpoints: Using apply_tesseract with reverse-mode AD (.backward(), torch.autograd.grad) requires the Tesseract to define a vector_jacobian_product endpoint. Forward-mode AD (torch.autograd.forward_ad) requires jacobian_vector_product.

  • torch.func transforms are not supported: apply_tesseract works with PyTorch's standard autograd API (.backward(), torch.autograd.grad, torch.autograd.forward_ad), but not with torch.func transforms (torch.func.vjp, torch.func.jvp, torch.func.grad, torch.func.vmap). These transforms create functionalized tensors that cannot be converted to NumPy arrays, which Tesseract endpoints require. Calling apply_tesseract inside a torch.func transform will raise a clear error.

License

Tesseract-Torch is licensed under the Apache License 2.0 and is free to use, modify, and distribute (under the terms of the license).

Tesseract is a registered trademark of Pasteur Labs, Inc. and may not be used without permission.

About

Execute + differentiate Tesseracts as part of PyTorch programs, with full support for reverse-mode and forward-mode AD. 🔦

Topics

Resources

Code of conduct

Contributing

Stars

Watchers

Forks

Releases

Contributors

Languages