Registration

autoform.extend.register_trace_type(type, aval_rule, /)[source]

Register a Python type as a traceable input type.

autoform.trace() treats registered Python types as dynamic leaves. During tracing, each concrete value is passed to aval_rule and replaced by an AVal that carries the abstract information needed by primitive rules.

Registering a trace type only teaches AutoForm how to abstract concrete inputs. It does not define concrete execution, AD behavior, batching, or Python operator syntax. Those are registered separately through primitive rules and the other helpers in this module.

Parameters:

• type (type) – Concrete Python type accepted as a dynamic input leaf.

• aval_rule (T) – Function from a concrete value to its abstract value.

Returns:

The registered rule.

Return type:

T

Example

>>> import functools as ft
>>> import autoform.extend as afe
>>> class Token: ...
>>> class TokenAVal(afe.AVal): ...
>>> @ft.partial(afe.register_trace_type, Token)
... def token_aval(value):
...     return TokenAVal()

autoform.extend.register_zero(aval_type, rule, /)[source]

Register the concrete zero value for an abstract value type.

Reverse-mode transforms use Zero to represent a missing or blocked cotangent. autoform.extend.materialize() later turns that symbolic zero into a concrete runtime value by looking up this rule.

Register this for differentiable value domains whose cotangents may need to flow through primitives that materialize zeros, such as autoform.pushforward() or custom pullback rules.

Parameters:

• aval_type (type[AVal]) – Abstract value type this zero rule handles.

• rule (T) – Function from an abstract value instance to a concrete zero.

Returns:

The registered rule.

Return type:

T

Example

>>> import functools as ft
>>> import autoform.extend as afe
>>> class TokenAVal(afe.AVal): ...
>>> @ft.partial(afe.register_zero, TokenAVal)
... def zero_token(aval):
...     return ""

autoform.extend.register_cotangent_accumulator(aval_type, rule, /)[source]

Register how to combine cotangents for an abstract value type.

autoform.pullback() can produce multiple cotangent contributions for the same input. AutoForm groups those contributions and calls the accumulator for the corresponding leaf AVal.

The rule receives the non-zero cotangents and the abstract value of the leaf being accumulated. Zeros are filtered before the rule is called.

Parameters:

• aval_type (type[AVal]) – Abstract value type this accumulator handles.

• rule (T) – Function of (cotangents, aval) returning one cotangent.

Returns:

The registered rule.

Return type:

T

Example

>>> import functools as ft
>>> import autoform.extend as afe
>>> class ScoreAVal(afe.AVal): ...
>>> @ft.partial(afe.register_cotangent_accumulator, ScoreAVal)
... def add_scores(cotangents, aval):
...     return sum(cotangents)

autoform.extend.register_non_dce(prim, /)[source]

Register a primitive as preserved during dead-code elimination.

Marks an extension primitive as semantically relevant even when its output is unused, such as a scoring, logging, or collection boundary.

Parameters:

prim (T) – Primitive preserved by autoform.dce().

Returns:

The registered primitive.

Return type:

T

autoform.extend.register_non_memoizable(prim, /)[source]

Register a primitive as excluded from autoform.memoize().

Marks an extension primitive as requiring repeated execution, such as stochastic sampling, scoring, logging, or calls that observe runtime state.

Parameters:

prim (T) – Primitive excluded from memoization.

Returns:

The registered primitive.

Return type:

T