Staging¶

Compile one callable signature into an immutable staged program.

Use the direct form with a callable, or omit function to create a decorator. Applying the decorator compiles the function and replaces it with a StagedProgram. Compilation traces the Python callable once with abstract values; later calls execute the graph without running or retracing the Python callable.

Declare the positional signature in exactly one of two ways: pass concrete examples to infer its array leaves, or pass specs containing ArraySpec and StaticSpec leaves. Keyword arguments have no example form and are declared with kw_specs in either case.

Parameters:

• function (Callable[..., Any] | None, default: None ) –

  Callable to compile. If omitted, return a decorator that compiles the callable it receives.

• *examples (Any, default: () ) –

  Concrete positional arguments whose pytree structure, shapes, dtypes, devices, and Python-scalar categories define the compiled signature. Wrap a non-array compile-time leaf in StaticSpec. Mutually exclusive with specs.

• specs (tuple[Any, ...] | None, default: None ) –

  Positional argument specification tree. Every leaf must be an ArraySpec or StaticSpec. Mutually exclusive with examples.

• kw_specs (dict[str, Any] | None, default: None ) –

  Mapping from keyword argument names to specification trees whose leaves are ArraySpec or StaticSpec. The mapping is combined with the positional signature declared by examples or specs.

• array_api_version (str | None, default: None ) –

  Array API revision to compile and store in the graph. With concrete examples and no explicit revision, Advect selects the newest supported revision served by their common array provider. With specs alone, it selects Advect's latest supported revision. An explicit revision must be supported by Advect and by the provider of every array example.

Returns:

• StagedProgram or callable –

  A fully compiled, single-signature StagedProgram when function is supplied; otherwise, a decorator that returns such a program. The program snapshots static inputs and captured constants, can be serialized, and never grows a polymorphic cache or retraces.

Raises:

• TypeError –

  If neither examples nor specs is supplied, if both are supplied, if an example is neither array-like nor a supported Python scalar nor wrapped in StaticSpec, if a specification contains another leaf type, or if the concrete array examples cannot use one common provider at the selected Array API revision.

• ValueError –

  If array_api_version is not a supported revision, or if abstract tracing finds incompatible shapes, dtypes, or operation semantics.

Examples:

Infer a direct-call signature from a concrete array:

>>> import advect as ad
>>> import numpy as np
>>> program = ad.stage(lambda x: x + 1, np.array([1.0, 2.0]))
>>> program(np.array([3.0, 4.0])).tolist()
[4.0, 5.0]

The decorator form uses explicit positional and keyword specifications:

>>> @ad.stage(
...     specs=(ad.ArraySpec((2,), "float32"),),
...     kw_specs={"scale": ad.ArraySpec((), "float32")},
... )
... def scale(x, *, scale):
...     return x * scale
>>> scale(
...     np.array([1.0, 2.0], dtype=np.float32),
...     scale=np.asarray(2.0, dtype=np.float32),
... ).tolist()
[2.0, 4.0]

Compile a reusable staged pullback program.

Parameters:

• f (StagedProgram) –

  Primal StagedProgram to transpose. Ordinary callables are not accepted; use vjp for a concrete dynamic pullback.

• argnums (int | tuple[int, ...] | None, default: None ) –

  Positional inputs to differentiate. An integer returns that input's gradient pytree directly, while a tuple returns a tuple in the given order. None selects input zero unless argnames is provided, in which case it selects no positional inputs. Negative indices are resolved against the program's positional signature.

• argnames (tuple[str, ...] | None, default: None ) –

  Keyword inputs from the staged signature to differentiate. Their gradients are returned in a dictionary keyed by name. When positional and named inputs are both selected, the result is (positional_gradients, named_gradients); the positional part follows the integer-versus-tuple rule above.

Returns:

• StagedProgram –

  An immutable, serializable program with the primal call signature plus a reserved keyword-only cotangent input. The cotangent has the primal output's pytree and leaf specifications. The program preserves the primal program's Array API revision.

Raises:

• IndexError –

  If a positional selection is out of range for the staged signature.

• TypeError –

  If f is not a StagedProgram, or a selected weak scalar signature is not real floating-point.

• ValueError –

  If positional selections are duplicated, a selected input is absent from the staged signature, or the primal signature already reserves the cotangent keyword.

• NoVJPError –

  If an operation on the differentiated path has no reverse-mode rule.

Examples:

>>> import advect as ad
>>> import numpy as np
>>> square = ad.stage(lambda x: x**2, np.array([1.0, 2.0]))
>>> pullback = ad.vjp_program(square)
>>> pullback(np.array([3.0, 4.0]), cotangent=np.ones(2)).tolist()
[6.0, 8.0]