mixle.models.streaming_transformer_leaf module

A streaming, non-buffering transformer-LM leaf – the keystone that removes the host-RAM materialization wall.

Where NeuralCategorical buffers the whole shard in the accumulator and news a fresh optimizer every M-step, this inverts the EM abstraction: the M-step OWNS a long-lived module + optimizer, and the accumulator’s seq_update IS one train step on a streamed micro-batch. value() returns (loss_sum, tokens) – two floats, telemetry only, NEVER the corpus – and estimate() is a no-op that wraps the live module.

This deliberately voids the sufficient-statistic algebra (value/combine are telemetry, not a foldable statistic): a sanctioned non-leaf carve-out (like NeuralGaussian.sample() raising), NOT an ABC change. It is the single-process prerequisite for the distributed (FSDP2) neural handle: each rank keeps its streamed shard resident and the only cross-rank collective becomes the in-backward gradient reduce-scatter, never a gather-suff-stats-to-root.

class StreamingTransformer(module, device='cpu')

Bases: SequenceEncodableProbabilityDistribution

Wraps a live, persistently-trained module. seq_log_density = next-token log p (eval/telemetry).

Parameters:

• module (Any)

• device (str)

classmethod from_config(vocab, *, d_model=128, n_layer=4, n_head=4, block=64, embedding=None, device='cpu')

Build the leaf from hyperparameters (no hand-built torch module) – the declarative estimator surface.

embedding optionally ties a shared CategoricalEmbedding across leaves.

Parameters:

• vocab (int)

• d_model (int)

• n_layer (int)

• n_head (int)

• block (int)

• embedding (Any)

• device (str)

Return type:

StreamingTransformer

log_density(xy)

Return the log-density or log-mass at a single observation.

Parameters:

xy (Any)

Return type:

float

predict(x)

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

StreamingTransformerSampler

seq_log_density(enc)

Return vectorized log-density values for sequence-encoded observations.

Parameters:

enc (Any)

Return type:

ndarray

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

StreamingTransformerEstimator

dist_to_encoder()

Return the data encoder used by this distribution for vectorized methods.

Return type:

StreamingTokenEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

StreamingTransformer

class StreamingTransformerSampler(dist, seed=None)

Bases: DistributionSampler

Parameters:

• dist (StreamingTransformer)

• seed (int | None)

sample(size=None, *, batched=True)

Draw observations.

Combinator samplers (mixture/sequence/…) accept batched. With batched=True (the default) each child stream is drawn in one vectorized call instead of a per-draw Python loop – far faster. Because every child sampler owns an independent RandomState, batching consumes each stream in the same order as the loop, so the draws are identical to the legacy path. batched=False forces that legacy per-draw loop as a guaranteed- stable reference. Leaf samplers are already vectorized and ignore the flag.

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

class StreamingTokenEncoder

Bases: DataSequenceEncoder

seq_encode(data)

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

data (list)

Return type:

tuple[ndarray, ndarray]

class StreamingTransformerAccumulator(module, lr, device)

Bases: SequenceEncodableStatisticAccumulator

seq_update = ONE train step against the PERSISTENT module + optimizer; value() = telemetry only.

Parameters:

• module (Any)

• lr (float)

• device (str)

seq_update(enc, weights, estimate)

Parameters:

• enc (Any)

• weights (Any)

• estimate (Any)

Return type:

None

update(x, weight, estimate)

Parameters:

• x (Any)

• weight (float)

• estimate (Any)

Return type:

None

initialize(x, weight, rng)

Parameters:

• x (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

Parameters:

• enc (Any)

• weights (Any)

• rng (Any)

Return type:

None

combine(other)

Parameters:

other (Any)

Return type:

StreamingTransformerAccumulator

value()

Return type:

tuple[float, int]

from_value(v)

Parameters:

v (tuple)

Return type:

StreamingTransformerAccumulator

acc_to_encoder()

Return type:

StreamingTokenEncoder

class StreamingTransformerAccumulatorFactory(module, lr, device)

Bases: StatisticAccumulatorFactory

Parameters:

• module (Any)

• lr (float)

• device (str)

make()

Return type:

StreamingTransformerAccumulator

class StreamingTransformerEstimator(module, lr=3e-3, device='cpu')

Bases: ParameterEstimator

Parameters:

• module (Any)

• lr (float)

• device (str)

accumulator_factory()

Return type:

StreamingTransformerAccumulatorFactory

estimate(nobs, suff_stat)

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

StreamingTransformer

class TransformerLMEstimator(vocab, *, d_model=128, n_layer=4, n_head=4, block=64, embedding=None, lr=3e-3, device='cpu')

Bases: StreamingTransformerEstimator

A Transformer language model as a fit-ready estimator: TransformerLMEstimator(vocab, d_model=..., ...).

The clean, declarative surface – no hand-built torch module, no Leaf(...).estimator() two-step. Drops into MixtureEstimator/CompositeEstimator like any other *Estimator. embedding optionally ties a shared CategoricalEmbedding (e.g. one word embedding across a mixture’s experts). TransformerLMEstimator(V, embedding=emb) and StreamingTransformer.from_config(V, embedding=emb).estimator() build the same thing.

Parameters:

• vocab (int)

• d_model (int)

• n_layer (int)

• n_head (int)

• block (int)

• embedding (Any)

• lr (float)

• device (str)

stream_fit(module, token_source, *, lr=3e-3, device='cpu', report_every=200, log=None)

Train module by streaming micro-batches from token_source (a generator). The accumulator holds the PERSISTENT optimizer and trains incrementally; its payload stays (loss_sum, tokens) – the corpus is never buffered. Returns (StreamingTransformer, (loss_sum, tokens)).

Parameters:

• module (Any)

• token_source (Any)

• lr (float)

• device (str)

• report_every (int)

• log (Any)

Return type:

tuple

StreamingTransformerLeaf

alias of StreamingTransformer

StreamingTransformerLeafEstimator

alias of StreamingTransformerEstimator