Use Control Flow Inside a Traced Function

autoform control-flow primitives keep branches and loops visible to the IR. Use them when the branch condition or loop state is part of the traced program.

Concept

Primitives · The IR · Pytrees

Route with switch

import autoform as af


def brief(text: str) -> str:
    return af.format("brief: {}", text)


def detailed(text: str) -> str:
    return af.format("detailed: {}", text)


branches = {
    "brief": af.trace(brief)("seed"),
    "detailed": af.trace(detailed)("seed"),
}


def route(kind: str, text: str) -> str:
    return af.switch(kind, branches, text)


ir = af.trace(route)("brief", "recursion")
print(ir.call("detailed", "recursion"))

Repeat with while_loop

import optree
import autoform as af


@optree.dataclasses.dataclass(namespace=af.PYTREE_NAMESPACE)
class State:
    text: str
    status: str


def keep_going(state: State) -> bool:
    return state.status == "continue"


def add_step(state: State) -> State:
    text = af.concat(state.text, "!")
    return State(text=text, status="done")


example = State(text="go", status="continue")
cond_ir = af.trace(keep_going)(example)
body_ir = af.trace(add_step)(example)

result = af.while_loop(cond_ir, body_ir, example, max_iters=3)
print(result)

The loop state is a registered pytree, using Optree’s dataclass integration.

Block Feedback with stop_gradient

import autoform as af


def combine(locked: str, editable: str) -> str:
    locked = af.stop_gradient(locked)
    return af.format("{}\n{}", locked, editable)


ir = af.trace(combine)("terms:", "draft answer")
inputs = ("terms:", "draft answer")
output, (locked_feedback, editable_feedback) = af.pullback(ir).call(inputs, "make clearer")

print(output)
print(locked_feedback)
print(editable_feedback)

The forward value of locked is unchanged. Feedback for that input is blocked.

Force Ordering with depends

import autoform as af


def ordered(topic: str) -> str:
    audit = af.format("audit {}", topic)
    answer = af.format("answer {}", topic)
    # return answer through a barrier that also waits for audit
    return af.depends(answer, audit)


ir = af.trace(ordered)("recursion")
scheduled = af.sched(ir)
print(scheduled.call("recursion"))

Use depends when a result should not become available until another traced value has also been evaluated, even though the returned value does not consume it directly. It does not force the computation that produces the returned value to start after the dependencies; a scheduler may still run independent producers concurrently and place the depends barrier after them.