Finite Difference Gradients

This example demonstrates how to make any Tesseract differentiable using finite differences, without implementing analytical gradient code.

Warning

This feature is experimental and available in tesseract_core.runtime.experimental. The API may change in future releases.

When to use finite differences

Finite differences are useful when:

  • Prototyping: You want to quickly test gradient-based workflows before investing time in analytical derivatives

  • Complex nested schemas: Your inputs have deeply nested structures (like the mesh data in this example) that would be tedious to differentiate manually

  • Legacy code: You’re wrapping existing code where deriving gradients would require significant refactoring

  • Verification: You want to cross-check your analytical gradient implementation

Trade-offs

Approach

Pros

Cons

Auto-diff (e.g. JAX, PyTorch)

Accurate, efficient, automatic

Requires compatible framework

Analytical

Accurate, efficient

Requires manual implementation

Finite Differences

No manual work, works with any code

Less accurate, more function evaluations

The example: meshstats_finitediff

This example is a variant of the meshstats example that computes statistics on volumetric mesh data. Instead of hand-written Jacobian implementations, it uses the finite difference helpers.

Input schema

The input is a complex nested structure representing a volumetric mesh:

class VolumetricMeshData(BaseModel):
    n_points: int
    n_cells: int
    points: Differentiable[Array[(None, 3), Float64]]
    num_points_per_cell: Array[(None,), Float64]
    cell_connectivity: Array[(None,), Int32]
    cell_data: dict[str, Array[(None, None), Float64]]
    point_data: dict[str, Array[(None, None), Float64]]

Implementing gradients with finite differences

You can implement all three AD endpoints (jacobian, jacobian_vector_product, vector_jacobian_product) with just a few lines of code:

from tesseract_core.runtime.experimental import (
    finite_difference_jacobian,
    finite_difference_jvp,
    finite_difference_vjp,
)

def jacobian(inputs, jac_inputs, jac_outputs):
    return finite_difference_jacobian(
        apply, inputs, jac_inputs, jac_outputs,
        algorithm="central",
        eps=1e-6,
    )

def jacobian_vector_product(inputs, jvp_inputs, jvp_outputs, tangent_vector):
    return finite_difference_jvp(
        apply, inputs, jvp_inputs, jvp_outputs, tangent_vector,
        algorithm="central",
        eps=1e-6,
    )

def vector_jacobian_product(inputs, vjp_inputs, vjp_outputs, cotangent_vector):
    return finite_difference_vjp(
        apply, inputs, vjp_inputs, vjp_outputs, cotangent_vector,
        algorithm="central",
        eps=1e-6,
    )

Choosing an algorithm

The algorithm parameter controls the finite difference method:

Algorithm

Function Evaluations

Accuracy

Use case

"central"

2n

Second-order

Default choice

"forward"

n + 1

First-order

When evaluations are expensive

"stochastic"

O(√n)

Approximate

High-dimensional inputs (1000s+ parameters)

where n is the number of array elements in the input that we’re differentiating with respect to.

The "stochastic" algorithm uses Simultaneous Perturbation Stochastic Approximation (SPSA), which estimates the Jacobian using random perturbation directions. It requires only 2 function evaluations per sample regardless of input dimension, making it useful when the number of input dimensions is very large. The resulting gradient estimates are unbiased, which makes them suitable for use with first-order SGD-family optimizers. They may not be suitable for applications requiring precise Jacobian values (e.g., Newton-type methods).

Note

All three algorithms use a single scalar eps for the perturbation magnitude. This means accuracy may degrade for inputs where the magnitudes of Jacobian entries vary by many orders of magnitude (e.g., physical applications with unnormalized inputs). If this is a concern, consider normalizing your inputs or implementing analytical gradients.

Making the algorithm configurable

The example also shows how to expose the algorithm choice as an input parameter:

from typing import Literal

FDAlgorithm = Literal["central", "forward", "stochastic"]

class InputSchema(BaseModel):
    mesh: VolumetricMeshData
    fd_algorithm: FDAlgorithm = "central"

def jacobian(inputs, jac_inputs, jac_outputs):
    return finite_difference_jacobian(
        apply, inputs, jac_inputs, jac_outputs,
        algorithm=inputs.fd_algorithm,  # User can choose at runtime
        eps=1e-6,
    )

Full source code

tesseract_api.py
# Copyright 2025 Pasteur Labs. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

"""Example Tesseract demonstrating finite difference gradient computation on mesh data.

This example shows how to use finite_difference_jacobian, finite_difference_jvp,
and finite_difference_vjp to make a Tesseract with complex nested schemas
differentiable without writing explicit gradient code.

This is a variant of the meshstats example that uses finite differences instead
of hand-written Jacobian implementations.
"""

from typing import Any, Literal

import numpy as np
from pydantic import BaseModel, Field, model_validator

from tesseract_core.runtime import (
    Array,
    Differentiable,
    Float64,
    Int32,
    ShapeDType,
)
from tesseract_core.runtime.experimental import (
    finite_difference_jacobian,
    finite_difference_jvp,
    finite_difference_vjp,
)

#
# Schemas
#

FDAlgorithm = Literal["central", "forward", "stochastic"]


class VolumetricMeshData(BaseModel):
    """Mock mesh schema; shapes not validated."""

    n_points: int
    n_cells: int

    points: Differentiable[Array[(None, 3), Float64]]
    num_points_per_cell: Array[(None,), Float64]  # should have length == n_cells
    cell_connectivity: Array[(None,), Int32]  # length == sum(num_points_per_cell)

    cell_data: dict[str, Array[(None, None), Float64]]
    point_data: dict[str, Array[(None, None), Float64]]

    @model_validator(mode="after")
    def validate_num_points_per_cell(self):
        if not isinstance(self.num_points_per_cell, np.ndarray):
            return self
        if len(self.num_points_per_cell) != self.n_cells:
            raise ValueError(f"Length of num_points_per_cell must be {self.n_cells}")
        return self

    @model_validator(mode="after")
    def validate_cell_connectivity(self):
        if not isinstance(self.cell_connectivity, np.ndarray):
            return self
        expected_len = sum(self.num_points_per_cell)
        if len(self.cell_connectivity) != expected_len:
            raise ValueError(f"Length of cell_connectivity must be {expected_len}")
        return self


class InputSchema(BaseModel):
    mesh: VolumetricMeshData = Field(
        description="The mesh you want summary statistics of"
    )
    fd_algorithm: FDAlgorithm = Field(
        default="central",
        description=(
            "Finite difference algorithm to use for gradient computation. "
            "Options: 'central' (most accurate), 'forward' (faster), "
            "'stochastic' (SPSA, scales to high dimensions)."
        ),
    )


class SummaryStatistics(BaseModel):
    first_point_coordinates: Differentiable[Array[(3,), Float64]] = Field(
        description="Coordinates of the first point defined in the mesh."
    )
    barycenter: Differentiable[Array[(3,), Float64]] = Field(
        description="Mean of all the points defined in the input mesh."
    )


class OutputSchema(BaseModel):
    statistics: SummaryStatistics = Field(description="Summary statistics of the mesh.")


#
# Required endpoints
#


def apply(inputs: InputSchema) -> OutputSchema:
    points = inputs.mesh.points

    statistics = SummaryStatistics(
        first_point_coordinates=points[0],
        barycenter=points.mean(axis=0),
    )

    return OutputSchema(statistics=statistics)


#
# Optional endpoints
#


def abstract_eval(abstract_inputs):
    """Calculate output shape of apply from the shape of its inputs."""
    input_points_shapedtype = abstract_inputs.mesh.points

    # get dimension of vector space points live in from input
    dim = input_points_shapedtype.shape[1]
    dtype = input_points_shapedtype.dtype
    return {
        "statistics": {
            "first_point_coordinates": ShapeDType(shape=(dim,), dtype=dtype),
            "barycenter": ShapeDType(shape=(dim,), dtype=dtype),
        }
    }


def jacobian(
    inputs: InputSchema,
    jac_inputs: set[str],
    jac_outputs: set[str],
):
    """Compute the Jacobian using finite differences.

    This implementation uses the finite_difference_jacobian helper which automatically
    computes numerical derivatives by perturbing each input element.
    """
    return finite_difference_jacobian(
        apply,
        inputs,
        jac_inputs,
        jac_outputs,
        algorithm=inputs.fd_algorithm,
        eps=1e-6,
    )


def jacobian_vector_product(
    inputs: InputSchema,
    jvp_inputs: set[str],
    jvp_outputs: set[str],
    tangent_vector: dict[str, Any],
):
    """Compute the Jacobian-vector product using finite differences.

    The JVP computes J @ v efficiently using directional derivatives,
    without explicitly forming the full Jacobian matrix.
    """
    return finite_difference_jvp(
        apply,
        inputs,
        jvp_inputs,
        jvp_outputs,
        tangent_vector,
        algorithm=inputs.fd_algorithm,
        eps=1e-6,
    )


def vector_jacobian_product(
    inputs: InputSchema,
    vjp_inputs: set[str],
    vjp_outputs: set[str],
    cotangent_vector: dict[str, Any],
):
    """Compute the vector-Jacobian product using finite differences.

    The VJP computes v @ J which is useful for backpropagation.
    """
    return finite_difference_vjp(
        apply,
        inputs,
        vjp_inputs,
        vjp_outputs,
        cotangent_vector,
        algorithm=inputs.fd_algorithm,
        eps=1e-6,
    )
tesseract_config.yaml
name: meshstats_finitediff
version: "1.0.0"
description: |
  Example Tesseract demonstrating finite difference gradient computation on mesh data.
  This is a variant of the meshstats example that uses finite_difference_jacobian,
  finite_difference_jvp, and finite_difference_vjp instead of hand-written gradients.

build_config:
  package_data: []
  custom_build_steps: []

See also