docs/LOOPS_AND_RECURRENCE.md

Loops, Recurrence, and Sequential Models in nFlow

The Core Principle: DAGs Only, No Cycles

nFlow's graph IR is a directed acyclic graph (DAG). This is not a limitation — it is the correct model for deep learning computation graphs. PyTorch, TensorFlow, ONNX, and JAX all use DAGs. Cycles in a computation graph don't have a well-defined execution order and can't be differentiated or compiled.

What looks like a "loop" in RNNs is always one of two things:

1. Explicit unrolling — you wire T time steps manually (classical RNNs)
2. Scan/fold operation — a single node that maps over a sequence (modern approach)

---

Pattern 1: LSTM Cell (Single Step)

nn.lstm_cell — computes one time step of an LSTM:

• Inputs: x [B, input_size], h [B, hidden_size], c [B, hidden_size]
• Outputs: h_new [B, hidden_size], c_new [B, hidden_size]

To model a 3-step LSTM in nFlow:

input[0] → lstm_cell → h1, c1
             ↑ h0, c0
input[1] → lstm_cell → h2, c2
             ↑ h1, c1
input[2] → lstm_cell → h3, c3
             ↑ h2, c2

You wire 3 nn.lstm_cell nodes in sequence. The "loop" is explicit at the graph level. This is exactly how torch.nn.LSTMCell works — and it's perfectly expressible as a DAG.

When to use: Fixed-length sequences known at model-design time. Clean, simple, easy to debug. Good for teaching, small models, and architectures with specific step counts.

---

Pattern 2: Scan / Fold Over Sequence (Recommended)

nn.scan — maps a cell over a variable-length sequence:

• Inputs: xs [B, T, input_size], h0 [B, hidden_size], c0 [B, hidden_size]
• Outputs: hs [B, T, hidden_size], h_T [B, hidden_size], c_T [B, hidden_size]
• Attrs: cell: str — uid of the cell composite to scan over (e.g. "user.lstm_cell")

This is equivalent to JAX's lax.scan, Haskell's scanl, and how modern implementations of RWKV, Mamba, and Linear RNNs are expressed.

xs [B,T,D] → nn.scan(cell="user.lstm_cell") → hs [B,T,H]
              ↑ h0, c0                           h_T, c_T

One node represents the entire recurrence. This is how torch.nn.LSTM and Mamba's SSM scan are compiled internally. The scan node exports to:

• PyTorch: nn.LSTM(...) or a torch.func.scan / manual for-loop
• Keras: keras.layers.RNN(cell, return_sequences=True)
• ONNX: ONNX::Loop / ONNX::Scan operators

When to use: Variable-length sequences, production models, Mamba/RWKV/S4/linear attention. The scan node abstracts away the loop entirely.

---

Pattern 3: Custom SSM (Mamba-style)

For Mamba, S4, RWKV, and similar linear recurrences (h_t = A·h_{t-1} + B·x_t):

Build a UserComposite for the SSM kernel:

1. Create a user.mamba_ssm composite in the Custom Nodes panel
2. Wire the inner graph: dt, A, B, C, x → selective_scan → y
3. The selective scan is itself a nn.scan over a custom cell

The .nfl file carries the paper_url: "https://arxiv.org/abs/2312.00752" so collaborators always know the source architecture.

---

Pattern 4: Bidirectional RNNs

Wire two scan nodes:

xs → scan_forward  → hs_fwd [B,T,H]  ──┐
xs → scan_backward → hs_bwd [B,T,H]  ──┤→ concat → output

scan_backward can use the same cell with direction: "backward" attr.

---

Why Not Explicit Cycles?

Some graph editors allow "feedback edges" that create cycles. This is fundamentally broken for ML:

1. No well-defined topo order → can't run inference, can't compute shapes
2. Not differentiable as written (need time-unrolling to get gradients)
3. ONNX, PyTorch JIT, TensorRT don't support cyclic graphs — you'd need a special lowering pass that nFlow avoids by design

The DAG-based approach (explicit unrolling + scan) is used by every production ML framework and compiler. nFlow follows this industry standard.

---

Quick Reference

Architecture	nFlow Pattern
LSTM (fixed T steps)	T × nn.lstm_cell wired in sequence
LSTM (variable length)	1 × nn.scan(cell="user.lstm_cell")
GRU	nn.gru_cell or nn.scan(cell="user.gru_cell")
Mamba SSM	user.mamba_ssm composite with nn.scan inside
RWKV	user.rwkv_block composite with nn.scan inside
Transformer (no recurrence)	Standard DAG — MHA + FFN blocks
Bidirectional	Two parallel nn.scan nodes + concat
Temporal Conv (TCN)	Standard conv nodes — fully DAG-compatible

---

Implementing nn.scan in nFlow

The nn.scan primitive is registered in nflow-ops/src/prims.rs. It:

• Declares attrs: cell: str (the composite uid to scan), axis: int = 1
• Shape inference: reads the cell's declared ports_out, infers sequence output shape
• Codegen: emits nn.LSTM / nn.GRU for built-in cells; torch.func.scan or manual loop for custom cells
• Execution: the CPU evaluator calls infer_user_composite iteratively for each step

This is the recommended way to add any sequential model to nFlow.