ARCHITECTURE
A stateful sequence component with model-level tooling around it.
ESA keeps its recurrent sequence state inside the ESA path, while MLBricks handles backend routing, model composition, training, lifecycle operations and generation utilities around that component.
LAYER FLOW
| Stage | Role |
|---|---|
| Input hidden states | Sequence representations enter the ESA layer. |
| ESA sequence processing | The layer produces updated hidden representations and can carry recurrent state across prefill/decode calls. |
| Backend resolution | auto,nativeorpytorchcontrols the supported execution route. |
| Parent model | Normalization, FFN/residual choices, embeddings and LM head are composed byESAModelor another MLBricks parent. |
SEQUENCE EXECUTION
Training and generation use different public workflows.
Training and full-sequence inference use normal module calls. Recurrent generation usesprefillfollowed bydecode_step.ESAModel.prepare_generation()and related helpers prepare the model-level generation path without changing the public model structure.
COMPOSITION
When ESA is placed insideBricksor another parentnn.Module, preferESA(..., device=None)and move the complete parent model once with.to(device). This keeps mixers, embeddings, residual paths, FFNs and inputs on the same device.