# Copyright 2025 Pasteur Labs. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""JAX gradient endpoint helpers, used in the JAX recipe.
These functions remove the boilerplate from a JAX-backed Tesseract
``tesseract_api.py`` by providing one-line implementations of the
``apply``, ``jacobian``, ``jacobian_vector_product``,
``vector_jacobian_product`` and ``abstract_eval`` endpoints.
"""
from collections.abc import Callable, Hashable
from typing import Any
import equinox as eqx
import jax
import jax.numpy as jnp
import numpy as np
from pydantic import BaseModel
from tesseract_core.runtime.cuda.ipc import has_cuda_array_interface
from tesseract_core.runtime.tree_transforms import (
LRUCache,
filter_func,
flatten_with_paths,
set_at_path,
)
[docs]
def as_jax_array(x: Any) -> jax.Array:
"""Convert an array to JAX, using DLPack for GPU inputs.
Existing JAX arrays are returned unchanged. Host inputs use ``asarray``
because decoded NumPy inputs can be read-only and cannot export DLPack.
"""
if isinstance(x, jax.Array):
return x
if has_cuda_array_interface(x):
return jnp.from_dlpack(x)
return jnp.asarray(x)
[docs]
def as_jax_arrays(tree: Any) -> Any:
"""Convert array leaves in a pytree to JAX, preserving non-array leaves."""
is_array = lambda x: (
isinstance(x, np.ndarray | np.generic) or has_cuda_array_interface(x)
)
return jax.tree.map(lambda x: as_jax_array(x) if is_array(x) else x, tree)
_jax_vjp_cache: LRUCache | None = None
def _set_jax_vjp_cache_size(size: int) -> None:
"""Enable or resize the VJP residual cache, discarding existing entries.
Public entry point: :func:`tesseract_core.runtime.experimental.set_jax_vjp_cache_size`.
"""
global _jax_vjp_cache
_jax_vjp_cache = LRUCache(maxsize=size) if size > 0 else None
def _as_bits(x: jax.Array) -> tuple[jax.Array, ...]:
"""Reinterpret ``x`` as unsigned integers, so that comparing is bitwise."""
if jnp.issubdtype(x.dtype, jnp.complexfloating):
return (*_as_bits(jnp.real(x)), *_as_bits(jnp.imag(x)))
if x.dtype == jnp.bool_:
return (x.astype(jnp.uint8),)
return (jax.lax.bitcast_convert_type(x, jnp.dtype(f"uint{x.dtype.itemsize * 8}")),)
@jax.jit
def _bitwise_equal(a: tuple[jax.Array, ...], b: tuple[jax.Array, ...]) -> jax.Array:
"""Whether two sequences of device arrays hold identical bytes.
Matches the byte comparison of host leaves, so ``-0.0`` and ``0.0`` differ
and a NaN equals itself.
"""
result = jnp.bool_(True)
for x, y in zip(a, b, strict=True):
for xb, yb in zip(_as_bits(x), _as_bits(y), strict=True):
result &= jnp.all(xb == yb)
return result
def _as_words(bits: jax.Array) -> tuple[jax.Array, ...]:
"""Split the unsigned integers from :func:`_as_bits` into ``uint32`` words."""
if bits.dtype.itemsize == 8:
return ((bits >> 32).astype(jnp.uint32), bits.astype(jnp.uint32))
return (bits.astype(jnp.uint32),)
def _mix(x: jax.Array) -> jax.Array:
"""The MurmurHash3 finalizer, which spreads every input bit over the output."""
x ^= x >> 16
x *= jnp.uint32(0x85EBCA6B)
x ^= x >> 13
x *= jnp.uint32(0xC2B2AE35)
x ^= x >> 16
return x
@jax.jit
def _fingerprint_jit(arrays: tuple) -> jax.Array:
h = jnp.uint32(0)
for x in arrays:
for bits in _as_bits(x):
for words in _as_words(bits.ravel()):
# Salting each word with its position makes permutations differ.
salt = jax.lax.iota(jnp.uint32, words.size) * jnp.uint32(0x9E3779B9)
h = _mix(h + jnp.sum(_mix(words ^ salt), dtype=jnp.uint32))
return h
class _ArrayLeaves:
"""JAX arrays of a cache key that live on the same devices.
Jitted functions hash and compare their contents where they live, so that
neither copies them to the host. They hash by dtype, shape, devices and a
32-bit fingerprint of their bytes, so that keys with different contents
rarely share a hash. The fingerprint only narrows the search: the cache's
dict tells keys that share a hash apart with ``__eq__``, which compares
contents exactly with :func:`_bitwise_equal`.
"""
def __init__(self, arrays: tuple[jax.Array, ...]) -> None:
self.arrays = arrays
self._meta = tuple(
(a.dtype.str, a.shape, frozenset(a.devices())) for a in arrays
)
self._hash: int | None = None
def __hash__(self) -> int:
# A dict hashes the key on every lookup, and each fingerprint waits for
# the device, so it is computed once per key.
if self._hash is None:
self._hash = hash((self._meta, int(_fingerprint_jit(self.arrays))))
return self._hash
def __eq__(self, other: object) -> bool:
return (
isinstance(other, _ArrayLeaves)
and self._meta == other._meta
and bool(_bitwise_equal(self.arrays, other.arrays))
)
def _cache_key(tree: Any) -> Hashable:
"""Build an :class:`LRUCache` key from a pytree's structure and leaves.
Array leaves are converted with :func:`as_jax_arrays`, as the endpoints do
before computing, so that the key describes the arrays the computation
sees. Their bytes stay where they are: the arrays on each set of devices
form one :class:`_ArrayLeaves` item at the end of the key, which hashes and
compares all of them in one call, and each array leaves the index of its
item at its position. The dtype and shape of every array are part of that
item, so that leaves with identical bytes but different interpretations
(e.g. ``int64[4]`` vs ``int64[2,2]``) don't collide. Arrays on different
devices go to different items because one jitted call cannot read them all.
Non-array leaves contribute their type alongside their value; they must be
hashable.
The key is returned as a tuple rather than collapsed with :func:`hash`,
so that :class:`LRUCache`'s dict lookup compares it with ``__eq__``
instead of trusting the hash alone. Collapsing to an ``int`` made every
hash collision a silent wrong-gradient bug: CPython reserves ``-1`` as an
error sentinel, so ``hash(-1) == hash(-2)``, and a Tesseract taking an
integer parameter could serve the backward pass of a different input.
The per-leaf type tag covers the other direction, where values compare
equal across types: ``1 == 1.0 == True``, so a bare value would still
collide for a field typed as a union.
"""
leaves, treedef = jax.tree.flatten(as_jax_arrays(tree))
# jax.PyTreeDef's __hash__ collides on dicts with different keys, so we
# use its string form as the discriminator instead.
items: list = [str(treedef)]
group_index: dict[frozenset, int] = {}
groups: list[list[jax.Array]] = []
for leaf in leaves:
if isinstance(leaf, jax.Array):
devices = frozenset(leaf.devices())
if devices not in group_index:
group_index[devices] = len(groups)
groups.append([])
groups[group_index[devices]].append(leaf)
items.append((_ArrayLeaves, group_index[devices]))
else:
items.append((type(leaf), leaf))
items.extend(_ArrayLeaves(tuple(group)) for group in groups)
return tuple(items)
[docs]
def jax_apply(apply_jit: Callable, inputs: BaseModel) -> dict:
"""Run ``apply_jit`` and, if caching is enabled, populate the VJP cache.
``apply_jit`` is assumed to already be JIT-compiled (e.g. wrapped with
``@eqx.filter_jit``); this helper does not jit it. The user-facing
``apply`` endpoint may want to do pre/post-processing around the call,
so we cannot wrap it in a jit internally.
When :data:`_jax_vjp_cache` is set (see
:func:`tesseract_core.runtime.experimental.set_jax_vjp_cache_size`),
the forward pass is run via ``jax.vjp`` so the resulting backward
function can be stashed and reused by a later :func:`jax_vjp` call.
Array inputs are converted to JAX before calling ``apply_jit``.
"""
inputs_dict = as_jax_arrays(inputs.model_dump())
if _jax_vjp_cache is None:
return apply_jit(inputs_dict)
# Compute forward pass via jax.vjp to cache residuals for a potential
# subsequent vector_jacobian_product call. eqx.partition separates
# array (differentiable) from non-array outputs; has_aux tells jax.vjp
# to only differentiate through the array outputs.
#
# The inputs are partitioned the same way before being handed to jax.vjp,
# which differentiates every leaf it receives: a str flag or an optional
# None would otherwise abort the call outright ("is not a valid JAX
# type"), so merely enabling the cache would break a Tesseract that runs
# fine without it. Static leaves are threaded back in as constants,
# mirroring jax_abstract_eval below. A scalar ``Differentiable[Float32]``
# validates to a numpy scalar, which ``is_array`` accepts, so it stays
# dynamic and keeps its gradient.
dynamic_inputs, static_inputs = eqx.partition(inputs_dict, filter_spec=eqx.is_array)
def _apply_for_vjp(dynamic_inputs: dict) -> tuple:
out = apply_jit(eqx.combine(static_inputs, dynamic_inputs))
diff_out, static_out = eqx.partition(out, eqx.is_array)
return diff_out, static_out
diff_primals, vjp_func, static_primals = jax.vjp(
_apply_for_vjp, dynamic_inputs, has_aux=True
)
out = eqx.combine(diff_primals, static_primals)
cotangent_template = jax.tree.map(jnp.zeros_like, diff_primals)
_jax_vjp_cache.put(
(apply_jit, _cache_key(inputs_dict)), (vjp_func, cotangent_template)
)
return out
[docs]
def jax_vjp(
apply_jit: Callable,
inputs: BaseModel,
vjp_inputs: set[str],
vjp_outputs: set[str],
cotangent_vector: dict[str, Any],
) -> dict[str, Any]:
"""Compute the vector-Jacobian product.
Reuses the cached backward from a prior :func:`jax_apply` call when one is available (see
:func:`tesseract_core.runtime.experimental.set_jax_vjp_cache_size`); otherwise falls
through to a JIT-compiled ``jax.vjp`` evaluation. The JIT
compilation happens internally on the first miss for a given
(input shape/dtype, path subset) combination and is cached for reuse.
"""
inputs_dict = as_jax_arrays(inputs.model_dump())
cotangent_vector = as_jax_arrays(cotangent_vector)
# Use get (not pop) so the cached residuals can serve multiple sequential
# vjp calls on the same inputs -- for example, when tesseract-jax's
# value_and_grad is followed by jax.jacrev, which decomposes into many
# vjp calls per output basis vector.
if (
_jax_vjp_cache is not None
and (cached := _jax_vjp_cache.get((apply_jit, _cache_key(inputs_dict))))
is not None
):
vjp_func, cotangent_template = cached
full_cotangent = jax.tree.map(jnp.zeros_like, cotangent_template)
full_cotangent = set_at_path(full_cotangent, cotangent_vector)
(all_input_cotangents,) = vjp_func(full_cotangent)
return flatten_with_paths(all_input_cotangents, include_paths=vjp_inputs)
# Cache disabled or cache miss: fall back to original JIT-compiled path.
return _vjp_jit(
apply_jit,
inputs_dict,
tuple(vjp_inputs),
tuple(vjp_outputs),
cotangent_vector,
)
[docs]
def jax_jvp(
apply_jit: Callable,
inputs: BaseModel,
jvp_inputs: set[str],
jvp_outputs: set[str],
tangent_vector: dict[str, Any],
) -> dict[str, Any]:
"""Compute the Jacobian-vector product via :func:`jax.jvp`.
JIT compilation is applied internally and cached per
``(input shape/dtype, jvp_inputs, jvp_outputs)`` combination.
"""
return _jvp_jit(
apply_jit,
as_jax_arrays(inputs.model_dump()),
tuple(jvp_inputs),
tuple(jvp_outputs),
as_jax_arrays(tangent_vector),
)
[docs]
def jax_jacobian(
apply_jit: Callable,
inputs: BaseModel,
jac_inputs: set[str],
jac_outputs: set[str],
) -> dict[str, dict[str, Any]]:
"""Compute the Jacobian via :func:`jax.jacrev`.
JIT compilation is applied internally and cached per
``(input shape/dtype, jac_inputs, jac_outputs)`` combination.
"""
return _jac_jit(
apply_jit,
as_jax_arrays(inputs.model_dump()),
tuple(jac_inputs),
tuple(jac_outputs),
)
[docs]
def jax_abstract_eval(apply_jit: Callable, abstract_inputs: BaseModel) -> dict:
"""Calculate the output shape of ``apply_jit`` from the shape of its inputs."""
is_shapedtype_dict = lambda x: type(x) is dict and (x.keys() == {"shape", "dtype"})
is_shapedtype_struct = lambda x: isinstance(x, jax.ShapeDtypeStruct)
jaxified_inputs = jax.tree.map(
lambda x: jax.ShapeDtypeStruct(**x) if is_shapedtype_dict(x) else x,
abstract_inputs.model_dump(),
is_leaf=is_shapedtype_dict,
)
dynamic_inputs, static_inputs = eqx.partition(
jaxified_inputs, filter_spec=is_shapedtype_struct
)
def wrapped_apply(dynamic_inputs: dict) -> dict:
inputs = eqx.combine(static_inputs, dynamic_inputs)
return apply_jit(inputs)
jax_shapes = jax.eval_shape(wrapped_apply, dynamic_inputs)
return jax.tree.map(
lambda x: (
{"shape": x.shape, "dtype": str(x.dtype)} if is_shapedtype_struct(x) else x
),
jax_shapes,
is_leaf=is_shapedtype_struct,
)
#
# Internal jit-compiled fallbacks (used on cache miss or when caching is disabled).
#
@eqx.filter_jit
def _jac_jit(
apply_jit: Callable,
inputs: dict,
jac_inputs: tuple[str, ...],
jac_outputs: tuple[str, ...],
) -> dict:
filtered_apply = filter_func(apply_jit, inputs, jac_outputs)
return jax.jacrev(filtered_apply)(
flatten_with_paths(inputs, include_paths=jac_inputs)
)
@eqx.filter_jit
def _jvp_jit(
apply_jit: Callable,
inputs: dict,
jvp_inputs: tuple[str, ...],
jvp_outputs: tuple[str, ...],
tangent_vector: dict,
) -> dict:
filtered_apply = filter_func(apply_jit, inputs, jvp_outputs)
return jax.jvp(
filtered_apply,
[flatten_with_paths(inputs, include_paths=jvp_inputs)],
[tangent_vector],
)[1]
@eqx.filter_jit
def _vjp_jit(
apply_jit: Callable,
inputs: dict,
vjp_inputs: tuple[str, ...],
vjp_outputs: tuple[str, ...],
cotangent_vector: dict,
) -> dict:
filtered_apply = filter_func(apply_jit, inputs, vjp_outputs)
_, vjp_func = jax.vjp(
filtered_apply, flatten_with_paths(inputs, include_paths=vjp_inputs)
)
return vjp_func(cotangent_vector)[0]