mixle.experimental.selective_scan module

E5 (part 1): the selective-scan (S6 / Mamba) module – a third ContextMechanism (see mixle/experimental/context_spine.py) alongside E1’s SlidingWindowSpine, targeting long, smooth, low-curvature dependencies that don’t compress into a fixed local window. See notes/designs/E5.md for the full design: why input-DEPENDENT (selective) Delta, A, B, C – not S4’s fixed, input-independent recurrence – is the property this mechanism exists for, why mamba-ssm is not a realistic dependency on this machine (no CUDA toolkit), and the exact S4D-real initialization this module uses.

_scan_layer is the ONE S6 recurrence implementation (a literal sequential Python loop over T, v1 per the design note’s explicit scope decision – a chunked/parallel scan is documented future work, not attempted here); both SelectiveScan.step() and mixle.experimental.ssm_hybrid.HybridBlock’s SSM branch call it, so there is exactly one scan, not two.

class SelectiveScanState(h=<factory>, pos=0)[source]

Bases: object

Per-layer recurrent state h ((batch, d_inner, d_state), None until the first step) plus the running absolute position counter – the SSM analogue of SlidingWindowState’s KV cache, except the state is already fixed-size (no window/cache-length bookkeeping needed).

Parameters:
class SelectiveScan(vocab, *, d_model=32, d_state=16, n_layer=2, expand=2)[source]

Bases: Module

E5 baseline: the S6/Mamba selective scan as a ContextMechanism.

Block shape mirrors SlidingWindowSpine’s pre-norm-residual convention (ln1 -> mixer -> residual -> ln2 -> mlp -> residual, weight-tied head) so the two mechanisms differ only in what the “mixer” is (notes/designs/E5.md). d_inner = expand * d_model (Mamba’s convention); the recurrent state h is carried across step calls exactly like SlidingWindowState’s KV cache, and detach does h.detach() per layer – same TBPTT contract, no window/cache-length bookkeeping needed since the state is already fixed-size.

Parameters:
log_density(x, y)[source]

x, y: (n, T) long tensors. Returns -mean_per_position_nll for each of the n sequences, each scored independently (state re-initialized per row) – one non-streaming forward per row, computed by calling init_state + step once per row exactly as a length-T, single-chunk stream would, not a separately-written scoring path (notes/designs/E5.md, “GradLeaf citizenship”).

Parameters:
Return type:

Any