mixle.experimental.context_spine module

E1: the chunked-recurrent training spine every Track-E long-context mechanism plugs into.

See notes/designs/E1.md for the design decisions (RoPE over learned-absolute position embeddings, windowed-mask derivation, why detach_horizon only ever cuts the backward graph and never the forward one, and why this ships its own tiny TBPTT driver instead of routing through GradLeaf).

mixle.models.transformer.CausalLM and mixle.models.streaming_transformer_leaf.StreamingTransformer both train bounded, independent micro-batches with a learned position table capped at block tokens – neither carries state across calls. ContextMechanism is the minimal protocol that adds streaming + carried state + truncated-backprop-through-time (TBPTT) on top, without touching either of those modules. SlidingWindowSpine is the E1 baseline mechanism (Transformer-XL-style stop-gradient KV carry); E2-E6 differ only in what step’s carried state contains.

class ContextMechanism(*args, **kwargs)

Bases: Protocol

The substrate contract every Track-E long-context mechanism implements.

step is per-position teacher-forced (returns the mean loss over every position in the chunk, not just the last one – unlike CausalLM.forward, which returns only the last position’s logits).

init_state(batch_size, *, device='cpu')

A fresh state for batch_size independent streams (empty cache / zero memory).

Parameters:

• batch_size (int)

• device (str)

Return type:

Any

step(state, chunk)

chunk = (x, y), (batch, T) long tensors. Returns (new_state, mean_loss).

Parameters:

• state (Any)

• chunk (tuple[Any, Any])

Return type:

tuple[Any, Any]

detach(state)

Stop-gradient the carried state (cuts the TBPTT backward graph at this point).

Parameters:

state (Any)

Return type:

Any

class SlidingWindowState(cache_k=<factory>, cache_v=<factory>, pos=0)

Bases: object

Per-layer stop-gradient KV cache plus the running absolute position counter (see E1.md’s RoPE note).

Parameters:

• cache_k (list[Any])

• cache_v (list[Any])

• pos (int)

class SlidingWindowSpine(vocab, *, d_model=32, n_layer=2, n_head=2, window=64, cp_size=1)

Bases: Module

E1 baseline: sliding-window exact attention with a stop-gradient carried KV cache (Transformer-XL style).

window=None (or window >= `` any sequence length this mechanism will ever see) makes ``step compute ordinary full causal self-attention with no truncation – the “full-attention-equivalent” configuration notes/designs/E1.md uses as the acceptance baseline, computed by the exact same code path multi-chunk streaming uses (not a second, independently-written transformer).

cp_size (E8, notes/designs/E8.md): context-parallel window sharding. cp_size=1 (the default) is the original single-device path above, completely unchanged – byte-identical, not just numerically close, so E1’s existing behavior and tests are untouched. cp_size > 1 shards the current step’s KV axis (cache ++ chunk) across cp_size simulated ranks via mixle.utils.parallel.context_parallel_spine for the attention sub-step of every layer; nothing else (embeddings, LayerNorm, MLP, head, cache bookkeeping) changes, since those are all per-position and need no communication.

Parameters:

• vocab (int)

• d_model (int)

• n_layer (int)

• n_head (int)

• window (int | None)

• cp_size (int)

train_tbptt(mechanism, state, chunks, opt, *, detach_horizon=1)

Stream chunks through mechanism, TBPTT-training with the given optimizer.

Every detach_horizon chunks (or at end of stream, whichever comes first): backward the mean accumulated loss, step the optimizer, then mechanism.detach(state) to cut the graph before continuing. detach_horizon=1 is literal per-chunk stop-gradient (Transformer-XL); a horizon spanning the whole stream means no mid-stream detach happens at all (see notes/designs/E1.md). Returns {"losses": [float, ...], "state": final_state} – one loss per chunk, detached telemetry.

Parameters:

• mechanism (ContextMechanism)

• state (Any)

• chunks (Any)

• opt (Any)

• detach_horizon (int)

Return type:

dict[str, Any]