Tesseract-JAX v0.4.0 released

Announcements
1

Release v0.4.0

Action required

Python 3.10 is no longer supported; JAX ≥ 0.7.0 required (#207)

The minimum supported Python version is now 3.11, and the minimum JAX version is now 0.7.0. Python 3.10 is nearly end-of-life and JAX dropped support for it in 0.7.0.

To upgrade, make sure your environment uses Python 3.11+ and a recent JAX, then update Tesseract-JAX as usual:

pip install --upgrade "tesseract-jax"

CI now also tests against Python 3.14.

Highlights

Much faster Jacobians via the jacobian endpoint (#191)

jax.jacfwd, jax.jacrev, and jax.jacobian over an apply_tesseract call (as well as jax.vmap over tangents/cotangents with unbatched primals) now dispatch to the Tesseract jacobian endpoint and contract the materialized Jacobian in a single pass, instead of calling the jvp/vjp endpoint once per batch element.

No code changes are required — existing transformations automatically take the fast path when the Tesseract exposes a jacobian endpoint. In benchmarks this is 14x faster for a 100×10 Jacobian and ~140x faster for a 100×100 Jacobian.

This unlocks differentiable-systems applications such as least-squares solves and root finding with JAX-based optimizers, as well as sensitivity analysis:

import jax
from tesseract_jax import apply_tesseract

# `jacfwd` / `jacrev` now use the jacobian endpoint under the hood
jac = jax.jacfwd(lambda x: apply_tesseract(tess, {"a": x})["result"])(x0)

Other changes

  • Fixed random deadlocks caused by allocating JAX memory from within Python callbacks; callbacks now return NumPy arrays only (#177).
  • Bumped the lowest runtime-tested JAX version to 0.6.2 internally before the 0.7.0 requirement landed (#201).
  • Docs: pulled in the new navbar from tesseract-core (#183) and corrected the comment on the VJP non-differentiable input placeholder (it is NaN, not zero) (#190).

What’s Changed

Features

  • Use jacobian endpoint (#191)
  • (deps) [breaking] Drop support by Python 3.10, test Python 3.14 in CI, bump jax lb to 0.7.0 (#207)

Bug Fixes

  • Prevent random deadlocks by not allocating JAX memory from within callbacks (#177)

Refactor

  • Bump lowest supported JAX version to 0.6.2 (#201)

Documentation

  • Pull in new navbar from tesseract-core (#183)
  • Correct comment on VJP non-diff input placeholder (NaN, not zero) (#190)

Full diff: Comparing v0.3.0...v0.4.0 · pasteurlabs/tesseract-jax · GitHub