Primitives¶

Define one atomic operation from its concrete implementation.

The implementation must have fixed named parameters: positional-or-keyword and keyword-only parameters are supported, while positional-only parameters, *args, and **kwargs are rejected. Calls still follow the implementation's normal Python signature.

static_argnames removes complete named arguments from tracing and stores them as operation attributes. nondiff_argnames keeps complete arguments as dynamic operands but supplies None tangents and suppresses their transpose contributions. The two sets must be disjoint. Derivative rules receive all remaining dynamic array/scalar leaves flattened in implementation-parameter and pytree order.

With residual=True, the implementation must return advect.PrimitiveResult; callers still receive only its output. Rules are attached to the returned handle.

Parameters:

• function (Callable[CallP, ResultT] | None, default: None ) –

  Concrete implementation, when the decorator is applied directly.

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

  Operation identity without the internal custom. prefix. By default Advect uses the implementation's module and qualified name. Use a stable explicit name for serialized artifacts.

• static_argnames (tuple[str, ...], default: () ) –

  Complete implementation arguments treated as concrete configuration.

• nondiff_argnames (tuple[str, ...], default: () ) –

  Complete dynamic arguments excluded from differentiation.

• residual (bool, default: False ) –

  Whether the implementation returns an invocation-local PrimitiveResult for an exact transpose.

Returns:

• Primitive or callable –

  A callable authoring handle, or a decorator that creates one.

Examples:

>>> import advect as ad
>>> import numpy as np
>>> @ad.primitive(name="examples.cube")
... def cube(value):
...     return value**3
>>> @cube.def_abstract
... def cube_abstract(value):
...     return value.spec
>>> @cube.def_jvp
... def cube_jvp(output, primals, tangents):
...     del output
...     (value,), (tangent,) = primals, tangents
...     return np.zeros_like(value) if tangent is None else 3 * value**2 * tangent
>>> from advect.testing import check_primitive
>>> sample = np.array([2.0])
>>> check_primitive(
...     cube,
...     primals=(sample,),
...     check=("abstract", "jvp", "transpose", "nested", "stage"),
... )
>>> ad.grad(lambda value: np.sum(cube(value)))(sample).tolist()
[12.0]

Attach the primitive's abstract staging rule.

The rule has the implementation's fixed named parameters. Advect preserves each dynamic argument's pytree while replacing its array/scalar leaves with advect.AbstractValue; declared static arguments arrive unchanged. Return the concrete output pytree with advect.ArraySpec or AbstractValue leaves.

The function is returned unchanged so this method can be used as a decorator.

Attach fn(output, primals, tangents, **static_attrs) as the JVP.

output has the implementation's public output pytree. primals and tangents are flat tuples with one entry per dynamic array/scalar leaf, in implementation-parameter and pytree order. Tangents may be None for inactive leaves and are always None for leaves of a declared nondifferentiable argument. Static arguments are passed by name. Return a tangent with the output pytree.

Write the rule as traceable, real-linear code so Advect can transpose it structurally and differentiate it again. The function is returned unchanged for decorator use.

Attach an ordinary or exact-residual transpose rule.

Ordinary primitives receive (cotangent, primals, output, **static_attrs). A primitive declared with residual=True receives (cotangent, primals, output, residual, **static_attrs). cotangent and output have the public output pytree; primals is the same flattened dynamic-leaf tuple used by the JVP. Return a flat tuple with one contribution per dynamic leaf in that order. Advect suppresses contributions for declared nondifferentiable arguments.

A rule may accept the optional keyword-only active_input_indices=None and return None for inactive contributions to avoid unnecessary work. Add an explicit transpose only when structural transposition cannot express the correct real adjoint, when an exact residual is required, or when measurement justifies a direct rule. The function is returned unchanged for decorator use.