Context Managers¶

These context managers wrap trace-time or execution-time behavior without transforming an IR.

autoform.fold()[source]¶

Evaluate immediately within the context.

Inside af.trace(...), primitive calls normally build IR equations. A fold block instead runs primitive implementations while tracing and returns concrete values that can be embedded as literals in the surrounding IR. If a primitive inside the block depends on a dynamic traced value, an AssertionError is raised. Outside tracing, fold is a no-op.

Example

>>> import autoform as af
>>> def program(x):
...     with af.fold():
...         prefix = af.concat("hello", " ")
...     return af.concat(prefix, x)
>>> ir = af.trace(program)("seed")
>>> len(ir.ir_eqns)
1
>>> ir.call("world")
'hello world'

Fold is useful when a trace-time computation should decide ordinary Python control flow. Autoform cannot stage Python branches whose conditions depend on dynamic IR values; those conditions must be known while tracing. A folded computation runs immediately, so its concrete result can safely choose the branch that is traced into the IR.

Example

>>> def program(x):
...     with af.fold():
...         route = af.concat("priority", ": high")
...     if route == "priority: high":
...         return af.concat("yes: ", x)
...     return af.concat("no: ", x)
>>> ir = af.trace(program)("seed")
>>> ir.call("answer")
'yes: answer'

Return type:: Generator[None, None, None]

autoform.tag(*tags)[source]¶

Attach tags to equations at trace time.

Equations built inside nested tag blocks receive the tags from all active blocks. Equations built after a block exits do not receive that block’s tags.

Example

>>> import autoform as af
>>> def program(x):
...     with af.tag("outer"):
...         head = af.concat(x, "!")
...         with af.tag("inner"):
...             return af.concat(head, "?")
>>> ir = af.trace(program)("seed")
>>> ir.ir_eqns[0].tags == frozenset({"outer"})
True
>>> ir.ir_eqns[1].tags == frozenset({"outer", "inner"})
True

Parameters:: tags (Hashable)
Return type:: Generator[tuple[Hashable, …], None, None]

autoform.memoize()[source]¶

Cache primitive results within the context.

Example

>>> import autoform as af
>>> def program(x):
...     a = af.concat(x, "!")
...     b = af.concat(x, "!")  # same call, will be cached
...     return af.concat(a, b)
>>> ir = af.trace(program)("test")
>>> with af.memoize():
...     result = ir.call("hello")
>>> result
'hello!hello!'

Tracing a program with memoize will act as compile-time deduplication of identical primitive calls (including stochastic primitives like lm_call()). Non-memoizable primitives are not memoized.

Example

>>> def program(x):
...     with af.memoize():
...         a = af.concat(x, "!")
...         b = af.concat(x, "!")  # same call, will be cached
...         return a, b
>>> ir = af.trace(program)("test")
>>> len(ir.ir_eqns)
1

Return type:: Generator[None, None, None]

autoform.collect(*, collection)[source]¶

Collect checkpoint values produced during IR execution.

collect is an execution-time context. Trace the program first, then place collect around ir.call(...) or ir.acall(...). Values are appended when executed autoform.checkpoint() primitives run.

Example

>>> import autoform as af
>>> def program(x):
...     normalized = af.format("item: {}", x)
...     normalized = af.checkpoint(normalized, key="normalized", collection="debug")
...     return af.concat(normalized, "!")
>>> ir = af.trace(program)("test")
>>> with af.collect(collection="debug") as collected:
...     result = ir.call("alpha")
>>> result
'item: alpha!'
>>> collected["normalized"]
['item: alpha']

Transformed IR execution is also execution, so collection works there too.

Example

>>> batched = af.batch(ir)
>>> with af.collect(collection="debug") as collected:
...     result = batched.call(["alpha", "beta"])
>>> result
['item: alpha!', 'item: beta!']
>>> collected["normalized"]
['item: alpha', 'item: beta']

Do not wrap trace construction with collect. Tracing builds IR equations; it does not produce concrete runtime checkpoint values. Do not use collect inside the function being traced either; during tracing, dynamic values are placeholders, not runtime values.

Parameters:: collection (Hashable) – The collection to filter marked values by. If …, collect all values.
Yields:: A dict that maps keys to lists of collected values.
Return type:: Generator[Collected, None, None]

autoform.inject(*, collection, values)[source]¶

Inject values for checkpoints within the context.

Values are consumed from lists in encounter order for each key. A dictionary produced by autoform.collect() can be supplied here to reproduce or modify checkpointed intermediates.

Around ir.call(...) or ir.acall(...), inject performs runtime checkpoint substitution by replacing matching checkpoint values as the IR executes.

Example

>>> import autoform as af
>>> def program(x):
...     normalized = af.format("item: {}", x)
...     normalized = af.checkpoint(normalized, key="normalized", collection="cache")
...     return af.concat(normalized, "!")
>>> ir = af.trace(program)("test")
>>> with af.inject(collection="cache", values={"normalized": ["cached item"]}):
...     ir.call("alpha")
'cached item!'

Inside the function being traced, inject performs trace-time specialization. A matching checkpoint is replaced by the injected literal while the IR is being built, so later calls to the traced IR reuse that specialized value.

Example

>>> def program(x):
...     normalized = af.format("item: {}", x)
...     with af.inject(collection="cache", values={"normalized": ["cached item"]}):
...         normalized = af.checkpoint(normalized, key="normalized", collection="cache")
...     return af.concat(normalized, "!")
>>> ir = af.trace(program)("test")
>>> [ir_eqn.prim.name for ir_eqn in ir.ir_eqns]
['format', 'concat']
>>> ir.call("alpha")
'cached item!'
>>> ir.call("beta")
'cached item!'

Parameters:

collection (Hashable) – The collection to filter checkpoint locations by.
values (Collected) – Dictionary mapping checkpoint keys to lists of values to inject.

Yields:

None.

Return type:

Generator[None, None, None]

autoform.lm_client(client)[source]¶

Set the LM client for all lm primitives.

The client must expose .completion() and .acompletion() matching LiteLLM’s chat completion signature.

Acceptable clients include the default direct LiteLLM adapter, a configured litellm.Router, or any wrapper object that forwards those two methods while preserving the LiteLLM request and response shapes.

Example

>>> import autoform as af
>>> from litellm import Router  
>>> client = Router(  
...     model_list=[
...         dict(model_name="gpt-4", litellm_params=dict(model="gpt-5.5")),
...     ],
...     max_parallel_requests=10,
... )
>>> with af.lm_client(client):  
...     ir.call(inputs)

Parameters:: client (LMClient)
Return type:: Generator[LMClient, None, None]