Branching and Loops in forward()

The power of ForwardConstruct is that Python control flow becomes graph topology. An if statement compiles to a conditional edge. self.loop() compiles to a graph cycle. A for loop over a proxy attribute compiles to an Each fan-out. The framework discovers both arms of a branch by re-tracing forward() with alternate branch decisions.

if compiles to conditional edges

from neograph import ForwardConstruct, Node, compile, run
from pydantic import BaseModel

class CheckResult(BaseModel, frozen=True):
    confidence: float

class Result(BaseModel, frozen=True):
    analysis: str

class QualityGate(ForwardConstruct):
    check   = Node(outputs=CheckResult, prompt='check', model='fast')
    deep    = Node(outputs=Result, prompt='deep-analysis', model='reason')
    shallow = Node(outputs=Result, prompt='quick-scan', model='fast')

    def forward(self, topic):
        checked = self.check(topic)
        if checked.confidence > 0.8:
            return self.shallow(checked)
        else:
            return self.deep(checked)

graph = compile(QualityGate())

The compiled graph has three nodes and a conditional edge after check:

If check.confidence > 0.8 at runtime, the graph routes to shallow.
Otherwise, it routes to deep.

How re-tracing works

The tracer uses the same strategy as torch.fx:

First trace: all if branches take the True arm. The tracer records which nodes appear.
Re-trace for each branch: flip that branch to False, re-run forward(), record which nodes appear.
Diff: nodes unique to the true trace become the true arm; nodes unique to the false trace become the false arm; shared nodes are unconditional.

The result is a node list annotated with _BranchMeta that the compiler lowers to add_conditional_edges.

Limitations

Only comparisons against constants are supported in v1:

# Supported — proxy attribute compared to a constant
if checked.confidence > 0.8:
    ...

# Supported — equality check
if checked.status == "approved":
    ...

# Not yet supported — comparison between two proxy values
if checked.score_a > checked.score_b:
    ...

Maximum 8 branches per forward(). Beyond that, ConstructError is raised — extract sub-pipelines to reduce branch count.

for compiles to Each fan-out

Iterating over a proxy attribute compiles to an Each modifier on the loop body’s nodes:

from neograph import ForwardConstruct, Node, compile, run
from pydantic import BaseModel

class ClusterGroup(BaseModel, frozen=True):
    label: str
    claims: list[str]

class Clusters(BaseModel, frozen=True):
    groups: list[ClusterGroup]

class VerifyResult(BaseModel, frozen=True):
    label: str
    passed: bool

class FanOutPipeline(ForwardConstruct):
    discover = Node(outputs=Clusters, prompt='discover', model='fast')
    verify   = Node(outputs=VerifyResult, prompt='verify', model='reason')

    def forward(self, topic):
        clusters = self.discover(topic)
        for group in clusters.groups:
            self.verify(group)

graph = compile(FanOutPipeline())

During tracing, for group in clusters.groups enters loop mode. The tracer:

Yields a single proxy item (enough to trace the loop body once).
Records that verify was called inside the loop.
Attaches Each(over="discover.groups", key="label") to the verify node.

The compiled graph runs verify once per item in discover.groups, collecting results as a dict keyed by label.

self.loop() compiles to graph cycles

Python for and while loops in forward() are traced once — the loop body runs during tracing but doesn’t produce a graph cycle. This is the same limitation as torch.jit.trace and JAX’s tracing: the tracer sees one unrolled pass, not a loop.

For iterative patterns that need a real back-edge in the compiled graph, use self.loop():

from neograph import ForwardConstruct, Node, compile
from pydantic import BaseModel

class Draft(BaseModel, frozen=True):
    content: str
    score: float = 0.0

class ReviewResult(BaseModel, frozen=True):
    score: float
    feedback: str

class Writer(ForwardConstruct):
    draft  = Node(outputs=Draft, prompt='draft', model='fast')
    review = Node(outputs=ReviewResult, prompt='review', model='reason')
    revise = Node(outputs=Draft, prompt='revise', model='reason')

    def forward(self, topic):
        d = self.draft(topic)
        d = self.loop(
            body=[self.review, self.revise],
            when=lambda r: r.score < 0.8,
            max_iterations=5,
        )(d)
        return d

graph = compile(Writer())

self.loop() takes a list of node references, an exit condition, and an iteration cap. It builds a sub-construct with a Loop modifier, producing a real cycle in the compiled graph. The when= callable receives the loop body’s latest output; return True to continue looping, False to exit.

This maps 1:1 to the programmatic equivalent construct | Loop(when=..., max_iterations=...) and to the YAML spec’s loop: block. The explicit primitive is language-agnostic — it works the same way regardless of host language, unlike tracing Python loops.

Single-node self-loop

For a single node that refines its own output, use loop_when= on @node instead:

@node(outputs=Draft, loop_when=lambda d: d is None or d.score < 0.8, max_iterations=5)
def refine(draft: Draft) -> Draft: ...

See the Loop with loop_when section for the full pattern.

Example: quality gate with retry

Combine branching and sequential calls for a retry pattern:

class RetryPipeline(ForwardConstruct):
    analyze  = Node(outputs=Analysis, prompt='analyze', model='reason')
    validate = Node(outputs=Validation, prompt='validate', model='fast')
    fix      = Node(outputs=Analysis, prompt='fix-issues', model='reason')
    report   = Node(outputs=Report, prompt='report', model='fast')

    def forward(self, topic):
        result = self.analyze(topic)
        checked = self.validate(result)
        if checked.score < 0.7:
            result = self.fix(result)
        return self.report(result)

This compiles to: analyze -> validate -> (if score < 0.7: fix) -> report. The fix node only runs when the quality gate fails.

Example: fan-out in a for loop

Process each cluster independently, then continue with the collected results:

class ClusterAnalysis(ForwardConstruct):
    discover = Node(outputs=Clusters, prompt='discover', model='fast')
    verify   = Node(outputs=VerifyResult, prompt='verify', model='reason')
    score    = Node(outputs=ScoreResult, prompt='score', model='fast')
    report   = Node(outputs=Report, prompt='report', model='fast')

    def forward(self, topic):
        clusters = self.discover(topic)
        for group in clusters.groups:
            self.verify(group)
            self.score(group)
        return self.report(clusters)

Both verify and score get Each modifiers and run once per cluster group. report runs once after all fan-out branches complete.

try/except in v1

try/except blocks in forward() are valid Python and don’t break tracing, but they don’t compile to fallback graphs.

During tracing, node calls are symbolic — they never raise. The except block is dead code during tracing because proxy operations always succeed. Only the try body’s nodes appear in the compiled graph.

class WithTryCatch(ForwardConstruct):
    primary  = Node(outputs=Result, prompt='primary', model='reason')
    fallback = Node(outputs=Result, prompt='fallback', model='fast')

    def forward(self, topic):
        try:
            return self.primary(topic)
        except Exception:
            return self.fallback(topic)  # never reached during tracing

In this example, only primary appears in the compiled graph. fallback is never traced.

For retry/fallback patterns in v1, use the conditional branch approach instead:

def forward(self, topic):
    result = self.primary(topic)
    if result.failed:
        return self.fallback(topic)
    return result