Types¶

Abstract Values¶

class autoform.extend.AVal[source]¶

Base class for abstract values used by traced programs.

Abstract values carry trace-time information about runtime values. Extension domains subclass AVal to describe the information primitive abstract rules need, such as shape, dtype, schema, or other static metadata.

Example

>>> import autoform.extend as afe
>>> class ArrayAVal(afe.AVal):
...     def __init__(self, shape, dtype):
...         self.shape = shape
...         self.dtype = dtype

class autoform.extend.StrAVal[source]¶

Abstract value for str leaves.

Example

>>> import autoform as af
>>> ir = af.trace(lambda x: x)("x")
>>> (x,) = ir.in_ir_tree
>>> x.aval
StrAVal()

class autoform.extend.IntAVal[source]¶

Abstract value for int leaves.

Example

>>> import autoform as af
>>> ir = af.trace(lambda x: x)(1)
>>> (x,) = ir.in_ir_tree
>>> x.aval
IntAVal()

class autoform.extend.FloatAVal[source]¶

Abstract value for float leaves.

Example

>>> import autoform as af
>>> ir = af.trace(lambda x: x)(1.0)
>>> (x,) = ir.in_ir_tree
>>> x.aval
FloatAVal()

class autoform.extend.BoolAVal[source]¶

Abstract value for bool leaves.

Example

>>> import autoform as af
>>> ir = af.trace(lambda x: x)(True)
>>> (x,) = ir.in_ir_tree
>>> x.aval
BoolAVal()

Primitives¶

class autoform.extend.Prim(name)[source]¶

Primitive operation key used by interpreter rule registries.

A primitive has no behavior by itself. Runtime, abstract, batching, and AD behavior are attached by registering rules keyed by the Prim instance.

Parameters:: name (str) – The name of the primitive.

Example

>>> import autoform.extend as afe
>>> add = afe.Prim("add")

class autoform.extend.Zero(aval, /)[source]¶

Symbolic zero cotangent for an abstract value.

Zero keeps reverse-mode and pushforward rules from materializing a concrete zero until one is actually needed. Use materialize() to replace symbolic zeros with concrete values.

Example

>>> import autoform.extend as afe
>>> z = afe.Zero(afe.StrAVal())
>>> afe.materialize(z)
''

Parameters:: aval (T)

IR¶

class autoform.extend.IR(ir_eqns, in_ir_tree, out_ir_tree)[source]¶

A traced AutoForm program.

An IR contains the ordered equations produced by tracing, plus the input and output IR trees that describe how runtime arguments and results are structured. Extension transforms may construct new IR values when they rewrite or wrap a program.

Parameters:

ir_eqns (list[IREqn]) – Ordered primitive equations.
in_ir_tree (Tree) – Tree describing the runtime input structure.
out_ir_tree (Tree) – Tree describing the runtime output structure.

class autoform.extend.IREqn(prim, in_ir_tree, out_ir_tree, params=None, tags=frozenset({}))[source]¶

One primitive application inside an IR.

An equation records the primitive to execute, the IR-shaped input and output trees, static primitive parameters, and the tags active when the equation was traced. Calling bind() executes the primitive under those tags.

Parameters:

prim (Prim) – Primitive represented by this equation.
in_ir_tree (Tree) – Input tree containing IR variables and concrete literals.
out_ir_tree (Tree) – Output tree containing IR variables and concrete literals.
params (dict[str, Any] | None) – Static parameters passed to the primitive rule.
tags (frozenset[Hashable]) – Tags associated with this equation.

class autoform.extend.IRVar(*, aval, source=None)[source]¶

Symbolic variable stored in IR trees.

IRVar leaves stand for runtime values inside traced programs. Each variable carries an AVal describing its abstract value, and an optional source variable used by transforms that create rewritten IR.

Parameters:

aval (AVal) – Abstract value for the runtime value represented by this variable.
source (IRVar | None) – Optional original variable this one was derived from.

Interpreters¶

class autoform.extend.Interpreter[source]¶

Base class for runtime primitive interpreters.

Subclass Interpreter to build an execution-time extension context. A custom interpreter usually stores the current active_interpreter as its parent, overrides interpret and ainterpret to handle new primitives.