Transforms

An IR transform is a function with this shape:

        flowchart TD
    input_ir["IR"] --> transform["IR transform"]
    transform --> output_ir["IR"]
    

The output is another executable IR. That one property is what makes composition ordinary Python function composition.

Transform	Returned IR expects	Returned IR produces	Use when
batch	Batched leaves where in_axes=True; broadcast leaves where in_axes=False.	Batched outputs.	Run the same program over many examples.
pushforward	Original inputs plus input tangents.	Original output plus output tangent.	Push a proposed input change forward.
pullback	Original inputs plus feedback on the output.	Original output plus input feedback.	Turn output critique into prompt/input critique.
sched	The same inputs as ir.	The same output as ir.	Overlap independent equations during async execution.
dce	The same inputs as ir.	The selected output shape, with unused leaves removed or replaced.	Drop work that cannot affect the needed outputs.
weighted	The same inputs as ir.	(output, path_weight).	Score one concrete path with reached factor calls.

IR Transforms

batch

batch(ir, /, *, in_axes=True) -> IR

batch vectorizes an IR over one or more input leaves. in_axes is a bool pytree matching the input structure: True means batched, False means broadcast. Higher-order primitives can have their own batch rules too, including while_loop.[1]

batched = af.batch(ir)
outputs = batched.call(["DNA", "gravity", "recursion"])

pushforward

pushforward(ir, /) -> IR

pushforward builds a forward-mode-style IR. The transformed IR takes primals and tangents, then returns output primals and output tangents.

pf = af.pushforward(ir)
output, tangent = pf.call(("topic",), ("make it concrete",))

pullback

pullback(ir, /) -> IR

pullback builds a reverse-mode-style IR. In autoform, cotangents are text feedback. The transformed IR takes the original inputs plus an output cotangent, then returns the output and input cotangents.

pb = af.pullback(ir)
output, input_feedback = pb.call(("topic",), "too abstract")

sched

sched(ir, /, *, cond=None) -> IR

sched groups independent equations into parallel stages. The resulting IR can still run with .call(...), but .acall(...) is where concurrent stages become useful.

import asyncio

scheduled = af.sched(ir)
result = asyncio.run(scheduled.acall("topic"))

weighted

weighted(ir, /) -> IR

weighted turns an IR into a path scorer. The returned IR runs one concrete path and returns the original output plus the product of reached factor weights.

scored = af.weighted(ir)
output, path_weight = scored.call("topic", 0.8)

dce

dce(ir, /, *, out_used=None) -> IR

dce removes equations that do not contribute to the selected output leaves.

trimmed = af.dce(ir)
result = trimmed.call("topic")

Composition

Composition works because every transform returns an IR:

transformed = af.batch(af.pullback(ir))
outputs, (topic_hints,) = transformed.call((topics,), critiques)

There is no special combined mode. pullback returns an IR; batch consumes that IR.

Order still matters:

Expression	Meaning
batch(pullback(ir))	Run many independent pullback calls at once. Each input pairs with its own output feedback.
pullback(batch(ir))	Treat the whole batched function as the program receiving feedback. The cotangent matches the batched output.
batch(weighted(ir))	Score many candidate paths separately. The result contains one weight per candidate.
weighted(batch(ir))	Score one batched path. Reached factors across the batched execution multiply into one weight.

Non-Transforms

Some nearby public APIs are intentionally not IR -> IR:

• custom is a decorator on traceable user functions. It marks a function boundary and lets transforms consult custom rules at that boundary. See Custom Rules.

• memoize is a context manager. It caches primitive results within a with block. See Cache Repeated Computations with memoize.

• lm_client is a context manager. It changes provider routing during execution. See Configure LiteLLM Routing.

• collect and inject are context managers. They capture or replace checkpointed values during execution. See Intercepts.

• tag and fold are context managers. They alter trace-time annotation or trace-time evaluation. See Tags and Fold.

The IR transforms reshape the IR. custom changes rule lookup at a boundary. Contexts wrap trace-time or execute-time behavior. They are complementary axes.

Transform and Execution Axes

The transform axis and execution axis compose independently:

import asyncio

transformed = af.batch(af.pullback(ir))

sync_result = transformed.call((topics,), critiques)
async_result = asyncio.run(transformed.acall((topics,), critiques))

The original function was not written as async def. Async execution is chosen when running the transformed IR. See Trace, IR, Execute for the execution split.

---

[1]

The batched while_loop implementation keeps an independent state for each batch item. Each iteration checks the condition for live items, runs the body only for items still active, and transposes between a batched pytree and per-item states internally. This lets different batch items exit on different iterations while the whole loop remains bounded by max_iters.