Python API (`tesseract_core.runtime`)¶

`tesseract_core.runtime`¶

The tesseract_core.runtime Python module contains some useful utilities to create and extend tesseract_api.py files. It is available in all Tesseract container images.

class tesseract_core.runtime.Array(*args, **kwargs)[source]¶

Generic Pydantic type annotation for a multi-dimensional array with a fixed shape and dtype.

Arrays will be broadcasted to the expected shape and dtype during validation, but dimensions must match exactly.

Polymorphic dimensions are supported by using None in the shape tuple. To indicate a scalar, use an empty tuple.

Arrays of any shape and rank can be represented by using … (ellipsis) as the shape.

Example

>>> class MyModel(BaseModel):
...     int_array: Array[(2, 3), Int32]
...     float_array: Array[(None, 3), Float64]
...     scalar_int: Array[(), Int16]
...     any_shape_array: Array[..., Float32]

You can serialize to (and validate from) different array encodings.

>>> model = MyModel(
...     int_array=np.array([[1, 2, 3], [4, 5, 6]]),
...     float_array=np.array([[1.0, 2.0, 3.0]]),
...     scalar_int=np.int32(42),
...     any_shape_array=np.array([True, False, True]).reshape(1, 1, 3),
... )

>>> model.model_dump_json(context={"array_encoding": "json"})

>>> model.model_dump_json(context={"array_encoding": "base64"})

or to binref:

>>> model.model_dump_json(
...     context={
...         "array_encoding": "binref",
...         "base_dir": "path/to/base",
...         "max_file_size": 10**8,
...     }
... )

In the ‘binref’ case you have to provide a base_dir to save/load binary (.bin) files. The .bin file(s) are written to context[‘base_dir’] / f”{context[‘__binref_uuid’]}.bin”. The ‘__binref_uuid’ is considered an internal variable and should not be modified manually! You can set a ‘max_file_size’ for the binary files. When this file size (in bytes) is reached, a new __binref_uuid (i.e. a new .bin) is created to append array data to.

class tesseract_core.runtime.Differentiable[source]¶

Type annotation for a differentiable array.

Example

>>> class MyModel(BaseModel):
...     array: Differentiable[Array[(None, 3), Float64]]

pydantic model tesseract_core.runtime.ShapeDType[source]¶

Data structure describing an array’s shape and data type.

Show JSON schema

{
   "title": "ShapeDType",
   "description": "Data structure describing an array's shape and data type.",
   "type": "object",
   "properties": {
      "shape": {
         "items": {
            "type": "integer"
         },
         "title": "Shape",
         "type": "array"
      },
      "dtype": {
         "enum": [
            "float16",
            "float32",
            "float64",
            "int8",
            "int16",
            "int32",
            "int64",
            "bool",
            "uint8",
            "uint16",
            "uint32",
            "uint64",
            "complex64",
            "complex128"
         ],
         "title": "Dtype",
         "type": "string"
      }
   },
   "required": [
      "shape",
      "dtype"
   ]
}

Config:

extra: str = ignore

Fields:

dtype (Literal['float16', 'float32', 'float64', 'int8', 'int16', 'int32', 'int64', 'bool', 'uint8', 'uint16', 'uint32', 'uint64', 'complex64', 'complex128'])
shape (tuple[int, ...])

field dtype: Literal['float16', 'float32', 'float64', 'int8', 'int16', 'int32', 'int64', 'bool', 'uint8', 'uint16', 'uint32', 'uint64', 'complex64', 'complex128'] [Required]¶

field shape: tuple[int, ...] [Required]¶

classmethod from_array_type(obj)[source]¶

Create a ShapeDType from an array annotation.

Return type:: _AnnotatedAlias

class tesseract_core.runtime.Float16¶: Float16 scalar array. Equivalent to Array[(), Float16].

class tesseract_core.runtime.Float32¶: Float32 scalar array. Equivalent to Array[(), Float32].

class tesseract_core.runtime.Float64¶: Float64 scalar array. Equivalent to Array[(), Float64].

class tesseract_core.runtime.Int8¶: Int8 scalar array. Equivalent to Array[(), Int8].

class tesseract_core.runtime.Int16¶: Int16 scalar array. Equivalent to Array[(), Int16].

class tesseract_core.runtime.Int32¶: Int32 scalar array. Equivalent to Array[(), Int32].

class tesseract_core.runtime.Int64¶: Int64 scalar array. Equivalent to Array[(), Int64].

class tesseract_core.runtime.UInt8¶: UInt8 scalar array. Equivalent to Array[(), UInt8].

class tesseract_core.runtime.UInt16¶: UInt16 scalar array. Equivalent to Array[(), UInt16].

class tesseract_core.runtime.UInt32¶: UInt32 scalar array. Equivalent to Array[(), UInt32].

class tesseract_core.runtime.UInt64¶: UInt64 scalar array. Equivalent to Array[(), UInt64].

class tesseract_core.runtime.Bool¶: Bool scalar array. Equivalent to Array[(), Bool].

class tesseract_core.runtime.Complex64¶: Complex64 scalar array. Equivalent to Array[(), Complex64].

class tesseract_core.runtime.Complex128¶: Complex128 scalar array. Equivalent to Array[(), Complex128].

`tesseract_core.runtime.experimental`¶

The experimental namespace includes features that are under active development and may not be fully stable. Use these at your own risk, as APIs, behaviors, and implementations can change without notice, or be removed entirely in future releases.

class tesseract_core.runtime.experimental.InputPath(*args, **kwargs)¶

Path type for Tesseract input files.

Resolves relative paths against RuntimeConfig().input_path and validates that the resolved path stays within the input directory and exists.

Use Annotated[InputPath, AfterValidator(...)] to add custom validation that runs after path resolution (receives the absolute, resolved path).

class tesseract_core.runtime.experimental.LazySequence(keys, getter)¶

Lazy sequence type that loads items from a file handle on access.

This allows users to define a sequence of objects that are lazily loaded from a data source, and validated when accessed.

When used as a Pydantic annotation, lazy sequences accept either a list of objects or a glob pattern to load objects from a file path.

Example

>>> class MyModel(BaseModel):
...     objects: LazySequence[str]
>>> model = MyModel.model_validate({"objects": ["item1", "item2"]})
>>> model.objects[0]
'item1'
>>> model = MyModel.model_validate({"objects": "@/path/to/data/*.json"})
>>> model.objects[1]
'item2'

class tesseract_core.runtime.experimental.OutputPath(*args, **kwargs)¶

Path type for Tesseract output files.

Validates that paths exist inside RuntimeConfig().output_path and strips the prefix when serializing.

In-memory values are absolute paths. Use Annotated[OutputPath, AfterValidator(...)] to add custom validation; validators receive absolute paths.

class tesseract_core.runtime.experimental.TesseractReference(tesseract)¶: Allows passing a reference to another Tesseract as input.

tesseract_core.runtime.experimental.finite_difference_jacobian(apply_fn, inputs, jac_inputs, jac_outputs, *, algorithm='central', eps=0.0001, num_samples=None, seed=None)¶

Compute the Jacobian of a Tesseract apply function using finite differences.

This function provides a generic way to make any Tesseract differentiable by computing gradients numerically. It can be used directly as the implementation of a jacobian endpoint.

Parameters:

apply_fn (Callable) – The Tesseract’s apply function with signature apply(inputs) -> outputs.
inputs (BaseModel) – The input data at which to compute the Jacobian.
jac_inputs (set[str]) – Set of input paths to differentiate with respect to.
jac_outputs (set[str]) – Set of output paths to compute derivatives of.
algorithm (Literal['central', 'forward', 'stochastic']) – The finite difference algorithm to use. Options are "central" (central differences, most accurate, 2 evaluations per element), "forward" (forward differences, faster, 1 extra evaluation per element), or "stochastic" (SPSA algorithm, scales better to high-dimensional inputs).
eps (float) – Perturbation magnitude for finite differences.
num_samples (int | None) – Number of random samples for the stochastic algorithm. Only used when algorithm="stochastic". Defaults to max(10, sqrt(n)) where n is the total number of input elements, providing O(sqrt(n)) cost instead of O(n) for full finite differences.
seed (int | None) – Random seed for reproducibility (only used with algorithm="stochastic").

Return type:

dict[str, dict[str, TypeAliasType]]

Returns:

A nested dictionary with structure {output_path: {input_path: jacobian_array}}, where each jacobian_array has shape (*output_shape, *input_shape).

Example

In a Tesseract’s tesseract_api.py:

from tesseract_core.runtime.experimental import finite_difference_jacobian


def jacobian(
    inputs: InputSchema,
    jac_inputs: set[str],
    jac_outputs: set[str],
):
    return finite_difference_jacobian(
        apply, inputs, jac_inputs, jac_outputs
    )

Note

This function is experimental and its API may change in future releases. It is useful for prototyping or when analytical gradients are difficult to derive, but numerical differentiation is generally less accurate and more computationally expensive than analytical methods.

tesseract_core.runtime.experimental.finite_difference_jvp(apply_fn, inputs, jvp_inputs, jvp_outputs, tangent_vector, *, algorithm='central', eps=0.0001)¶

Compute the Jacobian-vector product (JVP) using finite differences.

The JVP computes J @ v where J is the Jacobian and v is the tangent vector. This is done efficiently using directional derivatives without computing the full Jacobian.

Note: The "stochastic" algorithm is treated as "central" for JVP computation. JVP naturally requires only O(1) function evaluations regardless of input dimension (2 for central, 1 for forward), so stochastic estimation provides no benefit.

Parameters:

apply_fn (Callable) – The Tesseract’s apply function with signature apply(inputs) -> outputs.
inputs (BaseModel) – The input data at which to compute the JVP.
jvp_inputs (set[str]) – Set of input paths to differentiate with respect to.
jvp_outputs (set[str]) – Set of output paths to compute derivatives of.
tangent_vector (dict[str, TypeAliasType]) – Dictionary mapping input paths to tangent arrays.
algorithm (Literal['central', 'forward', 'stochastic']) – The finite difference algorithm to use. Options are "central" (most accurate, default) or "forward" (faster). The "stochastic" option is accepted but treated as "central".
eps (float) – Perturbation magnitude.

Return type:

dict[str, TypeAliasType]

Returns:

Dictionary mapping output paths to JVP result arrays.

Example

In a Tesseract’s tesseract_api.py:

from tesseract_core.runtime.experimental import finite_difference_jvp


def jacobian_vector_product(
    inputs: InputSchema,
    jvp_inputs: set[str],
    jvp_outputs: set[str],
    tangent_vector: dict[str, Any],
):
    return finite_difference_jvp(
        apply, inputs, jvp_inputs, jvp_outputs, tangent_vector
    )

Note

This function is experimental and its API may change in future releases.

tesseract_core.runtime.experimental.finite_difference_vjp(apply_fn, inputs, vjp_inputs, vjp_outputs, cotangent_vector, *, algorithm='central', eps=0.0001, num_samples=None, seed=None)¶

Compute the vector-Jacobian product (VJP) using finite differences.

The VJP computes v @ J where J is the Jacobian and v is the cotangent vector.

Note: For "central" and "forward" algorithms, the VJP is computed by explicitly computing the Jacobian rows and contracting with the cotangent vector. This requires O(n_inputs) function evaluations, the same cost as computing the full Jacobian. For high-dimensional inputs, consider using algorithm="stochastic".

Parameters:

apply_fn (Callable) – The Tesseract’s apply function with signature apply(inputs) -> outputs.
inputs (BaseModel) – The input data at which to compute the VJP.
vjp_inputs (set[str]) – Set of input paths to differentiate with respect to.
vjp_outputs (set[str]) – Set of output paths to compute derivatives of.
cotangent_vector (dict[str, TypeAliasType]) – Dictionary mapping output paths to cotangent arrays.
algorithm (Literal['central', 'forward', 'stochastic']) – The finite difference algorithm to use. Options are "central" (most accurate), "forward" (faster), or "stochastic" (SPSA, better for high-dimensional inputs).
eps (float) – Perturbation magnitude.
num_samples (int | None) – Number of random samples for the stochastic algorithm. Only used when algorithm="stochastic". Defaults to max(10, sqrt(n)) where n is the total number of input elements.
seed (int | None) – Random seed for reproducibility (only used with algorithm="stochastic").

Return type:

dict[str, TypeAliasType]

Returns:

Dictionary mapping input paths to VJP result arrays.

Example

In a Tesseract’s tesseract_api.py:

from tesseract_core.runtime.experimental import finite_difference_vjp


def vector_jacobian_product(
    inputs: InputSchema,
    vjp_inputs: set[str],
    vjp_outputs: set[str],
    cotangent_vector: dict[str, Any],
):
    return finite_difference_vjp(
        apply, inputs, vjp_inputs, vjp_outputs, cotangent_vector
    )

Note

This function is experimental and its API may change in future releases.

tesseract_core.runtime.experimental.jacobian_from_jvp(jvp_fn, inputs, jac_inputs, jac_outputs)¶

Compute the Jacobian by calling JVP with one-hot tangent vectors.

Requires N calls to JVP, where N is the total number of input elements. Output shapes are inferred from the first JVP probe.

Parameters:

jvp_fn (Callable) – The api_module.jacobian_vector_product callable.
inputs (Any) – Validated InputSchema instance.
jac_inputs (set[str]) – set[str] of input path strings.
jac_outputs (set[str]) – set[str] of output path strings.

Returns:

{input_path: array}} where each array has shape (*output_shape, *input_shape).

Return type:

dict[str, dict[str, np.ndarray]] with structure {output_path

tesseract_core.runtime.experimental.jacobian_from_vjp(vjp_fn, eval_fn, inputs, jac_inputs, jac_outputs)¶

Compute the Jacobian by calling VJP with one-hot cotangent vectors.

Requires M calls to VJP, where M is the total number of output elements.

Parameters:

vjp_fn (Callable) – The api_module.vector_jacobian_product callable.
eval_fn (Callable) – Either api_module.apply or api_module.abstract_eval, used to determine output shapes and dtypes. abstract_eval is preferred as it avoids a full forward pass.
inputs (Any) – Validated InputSchema instance.
jac_inputs (set[str]) – set[str] of input path strings.
jac_outputs (set[str]) – set[str] of output path strings.

Returns:

{input_path: array}} where each array has shape (*output_shape, *input_shape).

Return type:

dict[str, dict[str, np.ndarray]] with structure {output_path

tesseract_core.runtime.experimental.jvp_from_jacobian(jacobian_fn, inputs, jvp_inputs, jvp_outputs, tangent_vector)¶

Compute JVP as J @ t using the full Jacobian.

Parameters:

jacobian_fn (Callable) – The api_module.jacobian callable.
inputs (Any) – Validated InputSchema instance.
jvp_inputs (set[str]) – set[str] of input path strings.
jvp_outputs (set[str]) – set[str] of output path strings.
tangent_vector (dict[str, Any]) – dict mapping input paths to tangent arrays.

Return type:

dict[str, Any]

Returns:

dict mapping output paths to JVP result arrays.

tesseract_core.runtime.experimental.log_artifact(local_path)¶

Log an artifact to the current run context.

Return type:: None

tesseract_core.runtime.experimental.log_metric(key, value, step=None)¶

Log a metric to the current run context.

Return type:: None

tesseract_core.runtime.experimental.log_parameter(key, value)¶

Log a parameter to the current run context.

Return type:: None

tesseract_core.runtime.experimental.require_file(file_path)¶

Designate a file which is required to be present at runtime.

Parameters:: file_path (str | Path) – Path to required file. Must be relative to input_path assigned in tesseract run.
Return type:: Path

tesseract_core.runtime.experimental.set_jax_vjp_cache_size(size)¶

Enable or resize the VJP residual cache, discarding existing entries.

Experimental. The cache is opt-in and disabled by default; its performance characteristics are problem-dependent (see below) and the interface may change.

When jax_apply and jax_vjp (from tesseract_core.runtime.jax_recipes) are called in sequence on the same inputs (the typical pattern under tesseract-jax’s custom_vjp, including under jax.value_and_grad and jax.jacrev), caching the backward function produced by jax_apply’s internal jax.vjp call lets the subsequent jax_vjp skip the redundant forward pass. The win scales with the cost of the forward pass and goes the wrong way on trivially-small ones, so benchmark your own workload before enabling.

Best fit:

tesseract-jax workflows that go through custom_vjp (every gradient evaluation runs apply() first via fwd, so the cache reliably hits on the subsequent vjp() – true even under plain jax.grad).
Manual apply() followed by multiple vjp() calls on the same inputs (e.g. iterative solvers, CG-style inverse problems).

Poor fit:

Very small models where the forward pass is microseconds: the cache machinery overhead exceeds the saved work.

Parameters:: size (int) – Number of cache slots. 0 (the default) disables caching entirely (both jax_apply and jax_vjp then bypass the cache machinery). 1 covers the standard apply -> vjp pattern. Increase for workflows that interleave multiple apply() calls before their corresponding vjp() calls.
Return type:: None

tesseract_core.runtime.experimental.vjp_from_jacobian(jacobian_fn, inputs, vjp_inputs, vjp_outputs, cotangent_vector)¶

Compute VJP as v^T @ J using the full Jacobian.

Parameters:

jacobian_fn (Callable) – The api_module.jacobian callable.
inputs (Any) – Validated InputSchema instance.
vjp_inputs (set[str]) – set[str] of input path strings to differentiate w.r.t.
vjp_outputs (set[str]) – set[str] of output path strings to differentiate.
cotangent_vector (dict[str, Any]) – dict mapping output paths to cotangent arrays.

Return type:

dict[str, Any]

Returns:

dict mapping input paths to gradient arrays.

`tesseract_core.runtime.jax_recipes`¶

The tesseract_core.runtime.jax_recipes module contains utilities for working with JAX in tesseract_api.py, including support for efficient vector-Jacobian products (VJPs) when using JAX models in Tesseract apps. Used by default when initializing Tesseracts via tesseract init --recipe jax.

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.

tesseract_core.runtime.jax_recipes.jax_abstract_eval(apply_jit, abstract_inputs)[source]¶

Calculate the output shape of apply_jit from the shape of its inputs.

Return type:: dict

tesseract_core.runtime.jax_recipes.jax_apply(apply_jit, inputs)[source]¶

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 _jax_vjp_cache is set (see 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 jax_vjp() call. Otherwise this is just apply_jit(inputs.model_dump()).

Return type:: dict

tesseract_core.runtime.jax_recipes.jax_jacobian(apply_jit, inputs, jac_inputs, jac_outputs)[source]¶

Compute the Jacobian via jax.jacrev().

JIT compilation is applied internally and cached per (input shape/dtype, jac_inputs, jac_outputs) combination.

Return type:: dict[str, dict[str, Any]]

tesseract_core.runtime.jax_recipes.jax_jvp(apply_jit, inputs, jvp_inputs, jvp_outputs, tangent_vector)[source]¶

Compute the Jacobian-vector product via jax.jvp().

JIT compilation is applied internally and cached per (input shape/dtype, jvp_inputs, jvp_outputs) combination.

Return type:: dict[str, Any]

tesseract_core.runtime.jax_recipes.jax_vjp(apply_jit, inputs, vjp_inputs, vjp_outputs, cotangent_vector)[source]¶

Compute the vector-Jacobian product.

Reuses the cached backward from a prior jax_apply() call when one is available (see 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.

Return type:: dict[str, Any]

Python API (tesseract_core.runtime)¶

tesseract_core.runtime¶

tesseract_core.runtime.experimental¶

tesseract_core.runtime.jax_recipes¶

Python API (`tesseract_core.runtime`)¶

`tesseract_core.runtime`¶

`tesseract_core.runtime.experimental`¶

`tesseract_core.runtime.jax_recipes`¶