Skip to main content

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

NeuroScript
Loading editor...
PyTorch Output
Loading editor...

Here's what's happening:

  • blocks = unroll(num_layers): declares a named aggregate called blocks
  • Inside, block = TransformerBlock(...) defines the binding for each iteration
  • In the graph, blocks references the aggregate — the compiler expands it into a sequential chain where each iteration's output feeds the next iteration's input
  • Each block is 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

NeuroScript
Loading editor...
PyTorch Output
Loading editor...

@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

NeuroScript
Loading editor...
PyTorch Output
Loading editor...

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 @static is 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 in graph:

Try It Yourself

Experiment with unroll:

  • Change num_layers to see how the expanded graph changes
  • Add @static to see the difference in generated code
  • Compose the aggregate with other pipeline steps
  • Click "Show Analysis" to see the unrolled structure