Unroll
The unroll construct expands repeated layers at compile time. Instead of writing the same neuron call N times, you write it once and let the compiler generate the repetitions. Each iteration gets its own independent weights by default.
Why "Unroll"?
The name is intentional — this is not a runtime loop. The compiler literally unrolls the construct into N separate graph nodes before code generation. The count must be a literal integer or a parameter with a default value, never a runtime value.
Named Unroll Syntax
Declare an unroll block in the context: section with a named aggregate:
Transformer Stack
N transformer blocks chained sequentially via named unroll
Here's what's happening:
blocks = unroll(num_layers):declares a named aggregate calledblocks- Inside,
block = TransformerBlock(...)defines the binding for each iteration - In the graph,
blocksreferences the aggregate — the compiler expands it into a sequential chain where each iteration's output feeds the next iteration's input - Each
blockis an independent module with its own weights
Weight Sharing with @static
To reuse the same weights across all iterations (like a Universal Transformer), add @static to the binding:
Weight Sharing
@static makes all iterations share one set of weights
@static means one instance, reused across all iterations. The generated code creates a single TransformerBlock in __init__ and calls it num_steps times in forward.
End-to-End Example: GPT-2
unroll composes naturally with other pipeline steps — the aggregate is just another node in the chain:
GPT-2 Small
Full model: embedding → unrolled transformer blocks → head
This replaces what would otherwise be 12 hand-written TransformerBlock calls. The blocks aggregate sits naturally between embed and ln_f in the pipeline.
Key Rules
- The count must be a literal integer or a parameter with a default value — it must be resolvable at compile time
- Each iteration gets independent weights unless
@staticis used - The aggregate name (e.g.,
blocks) references the entire unrolled group as a sequential chain - Unroll blocks are declared in
context:and referenced by name ingraph:
Try It Yourself
Experiment with unroll:
- Change
num_layersto see how the expanded graph changes - Add
@staticto see the difference in generated code - Compose the aggregate with other pipeline steps
- Click "Show Analysis" to see the unrolled structure