Dextiny

Dextiny defines compact symbolic architecture contracts for arbitrary JAX callables and verifies real implementations against them without materializing input arrays.

Dextiny is library-agnostic. References can use ordinary JAX functions and modules from any JAX-traceable library, without requiring Dextiny-specific model classes.

Requirements

• python 3.12 or newer
• jax 0.11 or newer

Reference Chain

import dextiny as dx
import jax.numpy as jnp


tokens = dx.AbstractArray(
    "B S",
    batch_axis=0,
    dtype="int32",
    range="0 V",
)

table = dx.AbstractArray("V D", dtype="float32", name="table")

reference = (
    tokens
    .operate(
        lambda indices, weight: weight[indices],
        table,
        name="lookup",
    )
    .operate(jnp.sin)
    .operate(lambda hidden: hidden + 1, name="offset")
)

compiled = reference.compile(B=2, S=8, V=32, D=16)
print(compiled.render_history())

.operate() is a fluent wrapper around dx.operation(). It accepts the same function, operands, name, and kwargs, and returns a new immutable chain. The two forms can be mixed:

reference = (
    tokens
    .operate(jnp.sin)
    >> dx.operation(lambda value: value + 1, name="offset")
)

AbstractModule remains available for explicitly constructing a reusable stage. A plain callable on the right side of >> is wrapped automatically.

Dextiny uses >> rather than > because Python interprets a > b > c as a chained comparison, not a left-to-right pipeline.

Verify a Real Callable

real_table = jnp.zeros((32, 16), dtype=jnp.float32)


def real_model(input_ids):
    return jnp.sin(real_table[input_ids]) + 1


assert compiled.verify(real_model)

report = compiled.report(real_model)
print(report)
assert report.valid

Verify the same reference across several symbolic-dimension configurations by passing one mapping or a list or tuple of mappings to bindings:

assert compiled.verify(real_model, bindings={"B": 2, "S": 128})

report = compiled.report(
    real_model,
    bindings=[
        {"B": 1, "S": 1},
        {"B": 2, "S": 128},
        {"B": 4, "S": 2048},
    ],
)

assert report.valid
for result in report.results:
    print(result.bindings, result.valid)

A single mapping produces the usual VerificationReport. A list or tuple produces MultiVerificationReport, whose valid value is true only when every configuration matches. Each configuration overrides the dimensions stored in the compiled baseline and is abstractly recompiled and retraced. This catches Python and JAX paths selected only for particular shapes without materializing the configured arrays.

The report summarizes all three verification contracts:

Level: strict

Checks:
  Output contract: MATCH
  Operation sequence: MATCH (7 traced ops)
  Computation graph: MATCH

Inspect the traced operations behind the summary:

print(compiled.render_operations())  # Compiled reference only.

print(report.render_operations())        # Both ref and orig.
print(report.render_operations("ref"))   # Compiled reference.
print(report.render_operations("orig"))  # Original target callable.

ref is the compiled symbolic reference chain. orig is the original callable passed to compiled.report(...). Each traced operation includes its JAX primitive name, array signatures, input producers, output identities, and operation parameters. Pass show_parameters=False for a shorter view.

Rendered headings use the full names Reference traced operations and Original traced operations; ref and orig are only the API selectors.

The raw operation records are available as compiled.operations, report.ref_operations, and report.orig_operations.

Render the internal computation graph when auditing the calculation inside a reference stage:

print(report.render(show_graph=True))

The compiled contract compares:

• output PyTree structure
• output shapes and dtypes
• ordered JAX operation sequence, including nested JAXPRs
• computation graph dependencies between inputs, parameters, intermediate values, and returned outputs

Computation-graph verification distinguishes implementations that contain the same ordered operations but connect them differently. For example, a reference requiring left + right does not accept an implementation that calculates left + left.

Localized Differences

Every failed VerificationReport can locate and classify all mismatches:

report = compiled.report(original)
differences = report.diff(context=2)
difference = differences[0]

print(difference.classification)
print(difference.stage_index, difference.stage_name)
print(difference.traced_op_index)
print(report.render_diff(context=2))

Differences are returned as an ordered tuple, so they can be indexed or iterated directly with report.diff()[index]. A valid report returns ().

context selects how many traced operations before and after each divergence appear in the side-by-side Reference and Original operation graph. It does not retrace either callable. The diff is derived from the traces already stored in the report.

Differences are classified as:

• operation: different or missing JAX operations
• parameter: the same operation with different axes, permutations, dimension numbers, or other parameters
• constant: different embedded literal values under strict verification
• dependency: an operation or return value consumes a different producer
• output: final output shapes or dtypes differ
• contract: PyTree structure, tracing, or operation input/output contracts differ

Each result exposes structured details, reference_operation, original_operation, reference_context, and original_context fields for programmatic tooling. report.render_diff() renders every difference; report.render_diff(index) renders only the selected item.

Optional Numerical Parity

Abstract graph verification and numerical parity answer different questions. For example, x + x and x * 2 have different operation graphs but should produce the same values. Check this explicitly on a small concrete sample:

reference = (
    dx.AbstractArray("B D", dtype="float32")
    .operate(lambda x: x + x)
)
compiled = reference.compile(B=2, D=4)

original = lambda x: x * 2

assert not compiled.verify(original)             # Different graph.
assert compiled.verify_numerically(original)     # Equivalent values.

report = compiled.numerical_report(
    original,
    rtol=1e-5,
    atol=1e-6,
)
print(report)

numerical_report() generates a deterministic concrete input from the compiled AbstractArray. Pass sample= to control the tested values:

sample = jnp.ones((2, 4), dtype=jnp.float32)
report = compiled.numerical_report(original, sample=sample)

Integer arrays use their declared half-open range when one is available; otherwise their generated values are zero. Floating-point and complex inputs use deterministic random values selected by seed.

References containing symbolic operands, such as weights, need the same concrete values used by the original implementation. Bind them by their unique name or by the AbstractArray object itself:

weight_spec = dx.AbstractArray("K D", name="weight")
reference = dx.AbstractArray("B K").operate(
    lambda x, weight: x @ weight,
    weight_spec,
)
compiled = reference.compile(B=2, K=3, D=4)

report = compiled.numerical_report(
    original_module,
    operands={weight_spec: original_module.weight},
)

The default max_elements=1_000_000 bounds the total primary inputs, reference operands, and target arguments involved in the check. Dextiny also rejects a traced intermediate array larger than this limit before concrete execution. Raise the limit explicitly only when the resulting allocation is intentional. The guard cannot count arrays already captured inside the target callable, such as the parameters of an already-materialized model.

Numerical parity tests one sample, so it is evidence rather than a proof for all inputs. It remains entirely separate from verify() and report(); ordinary Dextiny verification continues to use only abstract arrays and makes no concrete allocations.

Gradient Verification

Verify the reverse-mode derivative graph with respect to the compiled primary input:

assert compiled.verify_gradients(real_model)

gradient_report = compiled.gradient_report(real_model)
print(gradient_report)

VJP is the default because it represents the backward path used by training and supports functions defined with jax.custom_vjp. Forward-mode and combined verification are explicit:

compiled.verify_gradients(real_model, mode="jvp")
compiled.verify_gradients(real_model, mode="both")

gradient_report.jvp and gradient_report.vjp expose the individual detailed reports when those modes were requested. Derivative operations can be audited directly:

print(gradient_report.render_operations("vjp"))

Gradient verification remains abstract: Dextiny constructs tangent and cotangent jax.ShapeDtypeStruct trees and traces jax.jvp or jax.vjp without materializing arrays. Reference operands and additional target arguments are treated as structural context; differentiation is currently with respect to the primary input. JAX does not permit forward-mode JVP through a jax.custom_vjp function, so use the default VJP mode for those callables.

Control Flow

Dextiny recursively verifies nested JAXPRs produced by:

• jax.lax.scan
• jax.lax.cond and jax.lax.switch
• jax.lax.while_loop
• jax.custom_jvp and jax.custom_vjp
• jax.jit and the legacy jax.experimental.pjit
• jax.shard_map

Nested inputs retain their exact parent lineage. In particular, condition selectors are excluded from branch arguments, and while_loop condition constants, body constants, and carry values are mapped independently. Paths in detailed operation reports identify the nested region, such as cond[branch_0], while[condition], scan[body], and shard_map[body].

For an unknown primitive, Dextiny maps a nested JAXPR only when its inputs have an unambiguous one-to-one correspondence with all parent inputs. Ambiguous mappings raise an error instead of silently selecting trailing operands. Python callback objects stored by custom_jvp and custom_vjp are excluded from forward parameter comparison; their derivative computations are verified through verify_gradients().

Typed Operation Parameters

Primitive parameters are stored as immutable CanonicalParameter values, not preformatted strings. Dimension numbers, permutations, axes, mappings, sets, dtypes, enums, partition specs, meshes, slices, dataclasses, and array metadata retain typed structure during comparison:

operation = compiled.operations[0]
parameters = dict(operation.parameters)

assert parameters["dimension_numbers"] == dx.canonical_parameter(
    (((1,), (0,)), ((), ()))
)

Lists and tuples share a canonical sequence representation so harmless container-format changes do not alter verification. Mappings and sets are ordered structurally. Array-valued metadata remains structural by shape and dtype rather than contents. Human-readable formatting is applied only by render_operations() or CanonicalParameter.render().

Verification Coverage

Custom primitives, callbacks, FFI calls, and external kernel calls may expose only their input/output contract to JAX tracing. Dextiny marks these operations as opaque instead of implying that their internal computation was inspected:

report = compiled.report(original)

print(report.coverage.status)  # "FULL", "PARTIAL", or "UNAVAILABLE"
print(report.coverage.reference.fraction)
print(report.opaque_operations)

assert report.valid            # No visible graph mismatch was found.
assert report.fully_verified   # Also requires complete visibility.

compiled.coverage describes the reference trace before an original callable is supplied. Reports provide separate reference and original TraceCoverage records with inspectable counts and opaque operation lists. Detailed operation rendering annotates opaque calls with their reason.

Nested computations provided by jit, control-flow primitives, custom_jvp, custom_vjp, and shard_map remain inspectable because Dextiny recursively traces their nested JAXPRs. A matching opaque operation keeps report.valid true but makes report.fully_verified false.

Verification Levels

Verification is strict by default:

Level	Verification
contract	Output PyTree structure, shapes, and dtypes
graph	Contract, operation sequence, operation parameters, and value dependencies
strict	Graph plus embedded scalar literal values

compiled.verify(real_model)                    # level="strict"
compiled.verify(real_model, level="graph")
compiled.verify(real_model, level="contract")

Strict verification distinguishes value * 2.0 from value * 3.0, including literals inside nested JAXPRs. Model operands remain structural: array weights are compared by shape and dtype rather than contents. Express scalar weights as operation operands to mark them structural as well.

Use an explicit wildcard when a reference constant's value is intentionally irrelevant:

reference = inputs >> dx.operation(
    lambda value, scale: value * scale,
    dx.wildcard("SCALE", dtype="float32"),
    name="scaled",
)
compiled = reference.compile(B=2, D=16, SCALE=1)

assert compiled.verify(lambda value: value * 3.0)

The wildcard key is resolved from the bindings passed to compile, so the same symbol can drive shapes and wildcard operands:

reference = (
    dx.AbstractArray("B D", dtype="float32")
    >> dx.operation(
        lambda value, scale: value * scale,
        dx.wildcard("B", dtype="float32"),
    )
)
compiled = reference.compile(B=8, D=16)

Changing B now requires editing only compile(B=...). Use dtype= when the binding is an integer but the operation expects a floating-point scalar. A wildcard matches an embedded constant in the original callable, not a runtime input.

Disable operation-sequence and computation-graph comparison when only the input/output contract matters:

matches = compiled.verify(real_model, level="contract")

check_primitives=False remains a compatibility alias for level="contract".

compiled(real_model) is shorthand for compiled.verify(real_model).

Symbolic Dimensions

Shapes are space-separated integer expressions:

hidden = dx.AbstractArray("B S H*D", dtype="bfloat16")
resolved = hidden.resolve({"B": 2, "S": 128, "H": 8, "D": 64})

Supported expression operators are +, -, *, //, and %. Dimensions must resolve to positive integers.

Value ranges are half-open metadata. range="0 V" means [0, V). Dextiny validates that the symbolic bounds resolve consistently; abstract JAX tracing cannot prove the values of a runtime array.

What Compilation Means

.compile() resolves dimensions and uses jax.make_jaxpr(..., return_shape=True) for every reference stage. It creates jax.ShapeDtypeStruct inputs, not concrete arrays, and does not invoke XLA executable compilation.

Verification traces the target using the same abstract input. It checks that the target follows the reference operation sequence and computation graph, but abstract verification cannot prove numerical equality. Numerical parity requires concrete test data.

Additional target arguments are supported:

matches = compiled.verify(lambda x, scale: x * scale, scale)

For complex keyword arguments or model-specific output selection, use a small wrapper:

report = compiled.report(
    lambda input_ids: model(input_ids, attention_mask=None)[0]
)

License

Dextiny is distributed under the Apache License 2.0. See LICENSE.md.