tesseract-jax 0.5.0

Tesseract-JAX

Tesseract-JAX is a lightweight extension to Tesseract Core that makes Tesseracts look and feel like regular JAX primitives, and makes them jittable, differentiable, and composable.

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

The API of Tesseract-JAX consists of a single function, apply_tesseract(tesseract_client, inputs), which is fully traceable by JAX. This enables end-to-end autodifferentiation and JIT compilation of Tesseract-based pipelines:

@jax.jit
def vector_sum(x, y):
    res = apply_tesseract(vectoradd_tesseract, {"a": {"v": x}, "b": {"v": y}})
    return res["vector_add"]["result"].sum()

jax.grad(vector_sum)(x, y) # 🎉

Quick start

1. Install Tesseract-JAX and get the example Tesseracts:

   $ pip install tesseract-jax
   $ git clone https://github.com/pasteurlabs/tesseract-jax
   

2. Use a Tesseract as part of a JAX program via the JAX-native apply_tesseract function:

   import jax
   import jax.numpy as jnp
   from tesseract_core import Tesseract
   from tesseract_jax import apply_tesseract
   
   # Serve the Tesseract in its own process (its environment is built on first use)
   t = Tesseract.from_source("tesseract-jax/examples/simple/vectoradd_jax/tesseract_api.py")
   t.serve()
   
   # Run it with JAX
   x = jnp.ones((1000,))
   y = jnp.ones((1000,))
   
   def vector_sum(x, y):
       res = apply_tesseract(t, {"a": {"v": x}, "b": {"v": y}}, vmap_method="sequential")
       return res["vector_add"]["result"].sum()
   
   vector_sum(x, y) # success!
   
   # You can also use it with JAX transformations like JIT and grad
   vector_sum_jit = jax.jit(vector_sum)
   vector_sum_jit(x, y)
   
   vector_sum_grad = jax.grad(vector_sum)
   vector_sum_grad(x, y)
   
   # vmap requires an explicit vmap_method — "sequential" is safe but slow
   # while "auto_experimental" or "expand_dims" is more efficient for Tesseracts that support batching.
   # See https://docs.pasteurlabs.ai/projects/tesseract-jax/latest/content/vmap-methods.html
   vector_sum_vmap = jax.vmap(vector_sum)
   vector_sum_vmap(x.reshape(10, 100), y.reshape(10, 100))
   

3. To share or deploy the Tesseract, build the same folder into a container image (this step requires Docker) and replace from_source with from_image:

   $ tesseract build tesseract-jax/examples/simple/vectoradd_jax
   

   t = Tesseract.from_image("vectoradd_jax")
   

   Use Tesseract.from_url(...) to reach a Tesseract that is already running elsewhere. apply_tesseract works the same way with all three.

Sharp edges

• Additional required endpoints: Tesseract-JAX requires the abstract_eval Tesseract endpoint to be defined to enable JAX tracing and FFI dispatch. To run a Tesseract that has no abstract_eval endpoint, call it directly through the Tesseract client instead. Additionally, many gradient transformations like jax.grad require vector_jacobian_product to be defined.

• Non-array outputs come from abstract_eval: an OutputSchema field that is not an array, such as a str or a bool, cannot enter a traced computation. As such, non-array outputs returned by apply are ignored and their value is instead taken from abstract_eval. apply_tesseract warns when the two differ; pass check_static_outputs=False or set TESSERACT_JAX_CHECK_STATIC_OUTPUTS=0 to skip the comparison. A field whose value depends on the input values belongs in the schema as an array.

License

Tesseract-JAX 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.