Primitives

Primitives are operations that tracing records as IR equations.

LM

autoform.lm_call(messages, /, *, model)[source]

Calls a language model with the given messages and model name using LiteLLM.

Parameters:

• messages (Messages) – A list of message dictionaries, each containing ‘role’ and ‘content’ keys.

• model (str) – The model name or active client model alias to use (e.g., “gpt-5.5”).

Returns:

The content of the model’s response as a string.

Return type:

str

Use lm_client() to configure provider-specific settings like max_tokens.

Example

>>> import autoform as af
>>> def program(name: str) -> str:
...     greeting = af.format("Hello, {}!", name)
...     sys = dict(role="system", content="translate the greeting to Korean")
...     usr = dict(role="user", content=greeting)
...     greeting = af.lm_call([sys, usr], model="gpt-5.5")
...     return greeting
>>> ir = af.trace(program)("World") 
>>> result = ir.call("x0") 

Example with lm_client():

>>> import autoform as af
>>> from litellm import Router  
>>> params_1024 = dict(model="gpt-5.5", max_tokens=1024)
>>> params_512 = dict(model="gpt-5.5", max_tokens=512)
>>> model_list = [
...     dict(model_name="gpt-5.5-1024", litellm_params=params_1024),
...     dict(model_name="gpt-5.5-512", litellm_params=params_512),
... ]
>>> router = Router(model_list=model_list)  
>>> def program(text: str, model: str):
...     msg = [{"role": "user", "content": af.format("Explain {} in one line.", text)}]
...     answer = af.lm_call(msg, model=model)
...     return af.concat("Answer: ", answer)
>>> ir = af.trace(program)("topic", "model")
>>> model_names = ["gpt-5.5-1024", "gpt-5.5-512"]
>>> with af.lm_client(router):  
...     result = af.batch(ir, in_axes=(False, True)).call("AI", model_names)

autoform.lm_schema_call(messages, /, *, model, schema)[source]

Calls a language model with an autoform schema response format.

The schema tree is built from nodes such as autoform.Int, autoform.Enum, and the other schema nodes exported by autoform.

Example

>>> import autoform as af
>>> answer = {
...     "name": af.Str() @ af.Doc("Subject name."),
...     "kind": af.Enum("summary", "definition") @ af.Doc("Answer kind."),
...     "score": af.Float(min=0, max=1) @ af.Doc("Confidence score."),
... } @ af.Doc("Answer object.")

Example with a registered pytree:

>>> import optree
>>> import autoform as af
>>> @optree.dataclasses.dataclass(namespace=af.PYTREE_NAMESPACE)
... class Answer:
...     answer: float
...     reasoning: str
>>> schema = Answer(
...     answer=af.Float() @ af.Doc("The numeric answer."),
...     reasoning=af.Str() @ af.Doc("The reasoning behind the answer."),
... )
>>> msgs = [dict(role="user", content="1 + 1?")]
>>> output = af.lm_schema_call(  
...     msgs,
...     model="openai/gpt-5.5",
...     schema=schema,
... )
>>> output  
Answer(answer=2.0, reasoning='Adding 1 and 1 gives 2.')

Parameters:

• messages (Messages)

• model (str)

• schema (Any)

Return type:

Any

String

autoform.format(template, *args, **kwargs)[source]

Format a string template with positional and/or keyword arguments.

Example

>>> import autoform as af
>>> af.format("Hello, {}!", "World")
'Hello, World!'
>>> af.format("Hello, {name}!", name="World")
'Hello, World!'
>>> af.format("{0}, {name}!", "Hi", name="World")
'Hi, World!'

Parameters:

template (str)

Return type:

str

autoform.concat(*args)[source]

Concatenates multiple strings into a single string.

Parameters:

*args – A variable number of string arguments to concatenate.

Returns:

A single string that is the concatenation of all input strings.

Return type:

str

Example

>>> import autoform as af
>>> result = af.concat("Hello, ", "world", "!")
>>> print(result)
Hello, world!

autoform.match(a, b, /)[source]

Check if two strings are equal.

This is a traceable version of == that works correctly during tracing.

Parameters:

• a (str) – First string

• b (str) – Second string

Returns:

True if strings are equal, False otherwise.

Return type:

bool

Example

>>> import autoform as af
>>> af.match("yes", "yes")
True
>>> af.match("yes", "no")
False

Control Flow

autoform.stop_gradient(x, /)[source]

Stops the gradient flow through the input during backpropagation.

Parameters:

x (Tree) – The input tree (e.g., a string, number, or nested structure)

Returns:

The same input tree with gradients stopped.

Return type:

Tree

Example

>>> import autoform as af
>>> def ir(x, y):
...     stopped = af.stop_gradient(x)
...     return af.concat(stopped, y)
>>> ir = af.trace(ir)("a", "b")
>>> pb_ir = af.pullback(ir)
>>> _, (cotangent_x, cotangent_y) = pb_ir.call(("a", "b"), "grad")
>>> cotangent_x
Zero(StrAVal())
>>> cotangent_y
'grad'

autoform.switch(key, branches, *args, **kwargs)[source]

Select and execute one of multiple IR branches based on a string key.

Parameters:

• key (str) – String key selecting which branch to execute.

• branches (Branches) – Dict mapping string keys to IR irs, each with compatible input signature.

• *args – Positional arguments passed to the selected branch.

Returns:

Result of branches[key].call(*args)

Raises:

KeyError – If key is not in branches.

Return type:

Tree

Example

>>> import autoform as af
>>> branches = {
...     "zero": af.trace(lambda x: af.concat("zero: ", x))("X"),
...     "one": af.trace(lambda x: af.concat("one: ", x))("X"),
...     "two": af.trace(lambda x: af.concat("two: ", x))("X"),
... }
>>> def ir(key, x):
...     return af.switch(key, branches, x)
>>> ir = af.trace(ir)("one", "hello")
>>> ir.call("one", "hello")
'one: hello'
>>> ir.call("zero", "hello")
'zero: hello'

autoform.while_loop(cond_ir, body_ir, init_val, *, max_iters)[source]

Repeatedly apply body_ir while cond_ir returns True.

• Loop continues while cond_ir(state) returns True

• body_ir is applied each iteration

• Returns final state when cond_ir returns False or max_iters reached

Parameters:

• cond_ir (IR) – IR that returns bool. Loop continues while True.

• body_ir (IR) – IR that transforms state, f: State -> State

• init_val (Tree) – Initial state

• max_iters (int) – Maximum iterations

Returns:

Final state when cond_ir returns False or max_iters reached.

Return type:

Tree

Example

>>> import autoform as af
>>> def cond(x):
...     return af.match(x, "go")
>>> def body(x):
...     return "stop"
>>> cond_ir = af.trace(cond)("...")
>>> body_ir = af.trace(body)("...")
>>> result = af.while_loop(cond_ir, body_ir, "go", max_iters=10)
>>> result
'stop'

Scheduling

autoform.depends(value, /, *deps)[source]

Annotate that value depends on the evaluation of deps.

This primitive inserts an ordering barrier without changing the forward value. The equations that produce value and deps may still run in the same scheduling level; the barrier’s output is available only after value and all deps have been evaluated.

Parameters:

• value (T) – The main value to return.

• *deps – Values that value depends on.

Returns:

The original value, through a barrier that also depends on deps.

Return type:

T

Example

>>> import autoform as af
>>> def program(x):
...     a = af.format("First: {}", x)
...     b = af.format("Second: {}", x)
...     return af.depends(b, a)  # return b after a has also run

Intercepts

autoform.checkpoint(value, /, *, key, collection=None)[source]

Tag a value with a collection and key for later collection.

checkpoint marks a value with a collection and key (unique identifier) that can be collected by collect. It acts as an identity operation in normal execution.

Parameters:

• value (Tree) – the value to mark (returned unchanged).

• key (Hashable) – unique identifier within the collection namespace.

• collection (Hashable | None) – optional collection for filtering (e.g., “debug”, “cache”, “metrics”).

Returns:

the input value unchanged.

Return type:

Tree

Example

>>> import autoform as af
>>> def program(x):
...     prompt = af.checkpoint(af.format("Q: {}", x), key="prompt", collection="debug")
...     response = af.concat(prompt, " A: 42")
...     return af.checkpoint(response, key="response", collection="debug")
>>> ir = af.trace(program)("test")
>>> with af.collect(collection="debug") as collected:
...     result = ir.call("What is 6*7?")
>>> result
'Q: What is 6*7? A: 42'
>>> collected["prompt"]
['Q: What is 6*7?']

Trace Weight

autoform.factor(weight, /, *, name=None)[source]

Multiply the current path weight by weight.

factor is neutral during ordinary execution. When an IR is transformed with weighted, each reached factor contributes to the returned path weight.

Parameters:

• weight (float)

• name (Hashable | None)

Return type:

None