autoform

autoform

Trace once. Transform freely.

Composable function transformations for LM programs.

Think JAX, but for LM programs.

Quickstart - Transforms - Concurrency - Debugging - Docs

pip install git+https://github.com/ASEM000/autoform.git

Quickstart

import autoform as af

def explain(topic: str) -> str:
    prompt = af.format("Explain {} in one paragraph.", topic)
    msg = dict(role="user", content=prompt)
    return af.lm_call([msg], model="gpt-5.2")

ir = af.trace(explain)("...")  # capture structure, no execution

Now transform it:

# execute
output = ir.call("quantum entanglement")

# batch: n inputs
outputs = af.batch(ir).call(["DNA", "gravity", "recursion"])

# pushforward: propagate input perturbations forward
output, tangent = af.pushforward(ir).call(("quantum entanglement", "add more examples"))

# pullback: propagate output feedback backward
output, grad = af.pullback(ir).call(("quantum entanglement", "too technical"))

# compose: batched differentiation
topics = ["DNA", "gravity", "recursion"]
critiques = ["too technical", "too brief", "too abstract"]
outputs, hints = af.batch(af.pullback(ir)).call((topics, critiques))

The last line is the point: batch(pullback(ir)), transformations compose.

Transforms

Transform	What it does
batch	Vectorize over inputs
pushforward	Forward-mode AD
pullback	Reverse-mode AD
sched	Auto-concurrent execution

Concurrency

sched finds independent LM calls. acall runs them concurrently.

scheduled = af.sched(ir)
result = await scheduled.acall("input") # acall for async

Debugging

Checkpoint intermediate values. Substitute on re-execution.

def pipeline(x: str) -> str:
    msg1 = dict(role="user", content=x)
    step1 = af.lm_call([msg1], model="gpt-5.2")
    step1 = af.checkpoint(step1, key="step1", collection="debug")
    
    msg2 = dict(role="user", content=step1)
    step2 = af.lm_call([msg2], model="gpt-5.2")
    return step2

ir = af.trace(pipeline)("...")

# capture
with af.collect(collection="debug") as captured:
    result = ir.call("input")

# substitute step1 value
with af.inject(collection="debug", values=dict(step1=["modified"])):
    result = ir.call("input")

⚠️ Early development: API may change.