Operator Registration

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

Register tracing dispatch for + on traced values with this aval.

This treats + as staged syntax while tracing. The rule is called during tracing, not during normal execution, and should usually bind an AutoForm primitive that implements the operation.

Parameters:

  • aval_type (type[AVal]) – Abstract value type of the left traced operand.

  • rule (T) – Function called as rule(left, right) for left + right and rule(right, left) for reflected right + left.

Returns:

The registered rule.

Return type:

T

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

Register tracing dispatch for - on traced values with this aval.

This treats - as staged syntax while tracing. The rule should normally bind the primitive that represents subtraction for the extension domain.

Parameters:

  • aval_type (type[AVal]) – Abstract value type of the traced operand that dispatches.

  • rule (T) – Function called as rule(left, right) for left - right and rule(right, left) for reflected right - left.

Returns:

The registered rule.

Return type:

T

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

Register tracing dispatch for * on traced values with this aval.

This treats * as staged syntax while tracing, instead of evaluating the operation with Python.

Parameters:

  • aval_type (type[AVal]) – Abstract value type of the traced operand that dispatches.

  • rule (T) – Function called as rule(left, right) for left * right and rule(right, left) for reflected right * left.

Returns:

The registered rule.

Return type:

T

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

Register tracing dispatch for / on traced values with this aval.

This treats true division as staged syntax while tracing.

Parameters:

  • aval_type (type[AVal]) – Abstract value type of the traced operand that dispatches.

  • rule (T) – Function called as rule(left, right) for left / right and rule(right, left) for reflected right / left.

Returns:

The registered rule, so the helper can be used as a decorator.

Return type:

T

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

Register tracing dispatch for @ on traced values with this aval.

This treats matrix multiplication as staged syntax while tracing. It is intended for domains such as arrays, matrices, or tensors where matrix multiplication should stage a primitive into the IR.

Parameters:

  • aval_type (type[AVal]) – Abstract value type of the traced operand that dispatches.

  • rule (T) – Function called as rule(left, right) for left @ right and rule(right, left) for reflected right @ left.

Returns:

The registered rule, so the helper can be used as a decorator.

Return type:

T

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

Register tracing dispatch for == on traced values with this aval.

Python equality on a traced value cannot be evaluated concretely during tracing. Register this when equality should become a staged primitive.

Parameters:

  • aval_type (type[AVal]) – Abstract value type of the left traced operand.

  • rule (T) – Function called as rule(left, right) for left == right.

Returns:

The registered rule, so the helper can be used as a decorator.

Return type:

T