Handling differentiability

Every Tesseract defines its interface through Pydantic BaseModel classes (InputSchema and OutputSchema). These schemas describe the structure, shapes, and dtypes of all array fields, and crucially, which fields are differentiable.

The Differentiable[...] annotation

Fields wrapped with Differentiable[...] participate in automatic differentiation. Fields without it are treated as constants by JAX’s AD machinery, even if their values change between calls.

from pydantic import BaseModel
from tesseract_core.runtime import Array, Differentiable, Float32

class InputSchema(BaseModel):
    x: Differentiable[Array[(3,), Float32]]   # differentiable
    label: Array[(1,), Float32]               # non-differentiable

class OutputSchema(BaseModel):
    loss: Differentiable[Array[(), Float32]]  # differentiable
    metadata: Array[(4,), Float32]            # non-differentiable

Fields can be arbitrarily nested (dicts, lists, and nested models). Differentiable[...] applies per-leaf:

class InputSchema(BaseModel):
    params: dict[str, Differentiable[Array[(None,), Float32]]]  # all leaves differentiable
    config: dict[str, Array[(None,), Float32]]                  # all leaves non-differentiable

When apply_tesseract is called inside a JAX differentiation context, Tesseract-JAX inspects these annotations to determine which inputs to request tangents/cotangents for, and which outputs to return derivatives of.

Non-differentiable inputs

When an input is not marked as Differentiable[...] in the Tesseract schema, differentiating with respect to it raises a ValueError in both forward and reverse mode. If you see this error, it likely means you forgot to annotate an input as Differentiable[...], or you are accidentally including a non-differentiable input in your differentiation.

  • Forward mode (jax.jvp, jax.jacfwd): providing a non-symbolic-zero tangent for a non-differentiable input raises a ValueError.

  • Reverse mode (jax.vjp, jax.grad, jax.jacrev): requesting a gradient with respect to a non-differentiable input raises a ValueError.

In both modes, use one of these strategies to exclude non-differentiable inputs from gradient computation:

Strategies for handling non-differentiable inputs

Closure: capture non-differentiable inputs outside the differentiated function:

# "b" is non-differentiable according to the Tesseract schema
def loss_fn(a):
    c = apply_tesseract(tess, {"a": a, "b": b})["c"]  # b captured from outer scope
    return jnp.sum(c**2)

jax.grad(loss_fn)(a)  # ✅ only differentiates w.r.t. "a"

argnums: explicitly select which arguments to differentiate:

def loss_fn(a, b):
    c = apply_tesseract(tess, {"a": a, "b": b})["c"]
    return jnp.sum(c**2)

jax.grad(loss_fn, argnums=0)(a, b)  # ✅ only differentiates w.r.t. "a"

stop_gradient: apply jax.lax.stop_gradient to the non-differentiable input inside the function, before passing it to apply_tesseract. This converts it to a concrete value, so no tangent reaches the primitive boundary:

Warning

stop_gradient changes the mathematical result of differentiation. Only use it if you are confident that gradient contributions through that path are genuinely undesirable or negligible. It is not a safe no-op.

def loss_fn(a, b):
    b = jax.lax.stop_gradient(b)
    c = apply_tesseract(tess, {"a": a, "b": b})["c"]
    return jnp.sum(c**2)

jax.grad(loss_fn)(a, b)  # ✅ stop_gradient prevents b from being differentiated

Non-differentiable outputs

When an output is not marked as Differentiable[...] in the Tesseract schema, Tesseract-JAX makes the problem explicit rather than silently producing wrong gradients:

  • Forward mode (jax.jvp, jax.jacfwd): the tangent for the non-differentiable output is NaN, which propagates to any downstream computation that depends on it. A ValueError is not raised here because the JVP rule is executed before any post-processing (such as pop or stop_gradient) can discard the output.

  • Reverse mode (jax.vjp, jax.grad, jax.jacrev): passing any concrete value as the cotangent for a non-differentiable output raises a ValueError. Only a symbolic zero jax._src.ad_util.Zero is accepted. If you see this error, it most likely means you forgot to annotate an output as Differentiable[...] in the Tesseract schema.

In both modes, you can use one of these strategies to exclude or insulate the non-differentiable output from gradient computation:

Strategies for handling non-differentiable outputs

Pop: remove it from the return value before differentiation:

def f(inputs):
    res = apply_tesseract(tess, inputs)
    res.pop("nondiff_res")
    return res

has_aux: return it as an auxiliary value outside the differentiated pytree:

def f(inputs):
    res = apply_tesseract(tess, inputs)
    return res["result"], res["nondiff_res"]  # (differentiable outputs, aux)

primals, f_vjp, nondiff_res = jax.vjp(f, inputs, has_aux=True)

stop_gradient: keep it in the return value but block gradient flow through it. In forward mode this produces a zero tangent instead of NaN; in reverse mode it produces a symbolic zero cotangent so no error is raised:

Warning

stop_gradient changes the mathematical result of differentiation. Only use it if you are confident that gradient contributions through that path are genuinely undesirable or negligible. It is not a safe no-op.

def f(inputs):
    res = apply_tesseract(tess, inputs)
    res["nondiff_res"] = jax.lax.stop_gradient(res["nondiff_res"])
    return res

Tip

For complex pytrees with many mixed differentiable and non-differentiable leaves, equinox.partition provides a convenient way to split and recombine them.

Note that the cotangent/tangent pytree structure must always match the function’s output structure. If you exclude outputs via pop or has_aux, including them in the cotangent raises a ValueError:

ValueError: unexpected tree structure of argument to vjp function:
  got PyTreeDef({'nondiff_res': *, 'result': *}), but expected PyTreeDef({'result': *})

Non-array outputs

An OutputSchema may carry fields that are not arrays at all, such as a backend name or a convergence flag:

class OutputSchema(BaseModel):
    y: Differentiable[Array[(3,), Float64]]
    backend: str = "reference"
    converged: bool = True

apply_tesseract returns those fields next to the arrays. They are plain Python objects rather than tracers, so a jit-ed function can branch on them at trace time:

@jax.jit
def f(x):
    out = apply_tesseract(tess, {"x": x})
    scale = 1.0 if out["converged"] else 0.0
    return out["y"] * scale

A JAX primitive can only return arrays, so a non-array field never enters the computation. apply_tesseract takes its value from abstract_eval and puts it back into the output pytree afterwards. Whatever apply returns for that field is therefore ignored, and apply_tesseract warns when the two disagree:

UserWarning: Tesseract returned the static output ['backend'] as 'fallback' from
apply, but abstract_eval reported 'reference'. ...

A field whose value depends on the input values belongs in the schema as an array instead.

Turning the check off

Pass check_static_outputs=False to skip the comparison for one call, or set TESSERACT_JAX_CHECK_STATIC_OUTPUTS=0 to skip it for the whole program. The keyword argument wins when both are set:

out = apply_tesseract(tess, {"x": x}, check_static_outputs=False)

Skipping the check avoids building the keypaths the warning needs. The values the caller receives are the same either way.

Higher-order derivatives

Tesseracts only expose first derivatives (the jacobian, jacobian_vector_product and vector_jacobian_product endpoints), so a function that calls apply_tesseract can be differentiated once, but not twice:

def f(x):
    return apply_tesseract(tess, {"x": x})["y"]

jax.grad(f)(x)  # ✅ first derivative

jax.hessian(f)(x)
# RuntimeError: Cannot take higher-order derivatives of 'jacobian'

jax.grad(jax.grad(f))(x)
# RuntimeError: Cannot differentiate a Tesseract derivative endpoint with respect
# to its primal inputs, as this needs a second derivative.

Which of the two errors you see depends on how the second derivative is taken. “Primal inputs” means the values f is evaluated at, here x.

Differentiating with respect to a tangent / cotangent works, because a derivative is linear in its direction:

v = jnp.ones_like(x)
jax.grad(lambda v: jax.jvp(f, (x,), (v,))[1])(v)  # ✅

How Tesseract-JAX handles partial differentiation (under JIT, it’s automatic!)

When you differentiate with respect to only some of a Tesseract’s differentiable inputs, or use only some of its differentiable outputs, tesseract-jax narrows the derivative request so the Tesseract computes only the sub-block it needs. The jacobian, jacobian_vector_product and vector_jacobian_product endpoints are asked for the minimal set of input columns and output rows.

Two mechanisms drive this:

  • Trace-time restriction is always applied. If a tangent is a symbolic zero — because you differentiate with respect to a subset of arguments, or wrap an input in jax.lax.stop_gradient — the corresponding column is dropped at trace time. This holds whether or not the call is under jax.jit.

  • Dead-code elimination (DCE) prunes what survives downstream. When a differentiated input or output carries a tangent but its result is discarded afterwards (e.g. subscripting the output of jax.grad or jax.jacfwd), the request is narrowed only if JAX runs DCE. DCE runs under jax.jit (any mode) and for un-jitted reverse mode (jax.grad / jax.jacrev). It does not run for un-jitted forward mode: an eager jax.jacfwd (or jax.jvp) that uses only some outputs will still request every differentiable output row.

Output rows and input-gradient columns are both pruned this way:

def f(x, y):
    # `tess` returns {"a": ..., "b": ...}, both differentiable; only "a" is used.
    return apply_tesseract(tess, {"x": x, "y": y})["a"]

jax.jit(jax.jacfwd(f))(x, y)  # requests only the "a" rows
jax.jacrev(f)(x, y)           # requests only the "a" rows (reverse mode runs DCE)
jax.jacfwd(f)(x, y)           # ⚠️ requests "a" and "b" rows, then discards "b"

# Discarding an input gradient downstream prunes its column the same way:
g = lambda x, y: jax.grad(lambda i: apply_tesseract(tess, i)["a"].sum())(
    {"x": x, "y": y}
)["x"]
jax.jit(g)(x, y)  # requests only the gradient w.r.t. "x"
g(x, y)           # ⚠️ requests both, then discards "y" (no DCE eagerly)

Wrapping the outer call in jax.jit recovers the pruning in the last case of each block. This only affects how much the Tesseract computes, never the returned value.