mixle.models package

Applied models – richer, domain-specialized families that plug into the same contract as mixle.stats.

Where mixle.stats holds the elementary distributions (a Gaussian, a Poisson, a categorical), this package holds the models that are more than one elementary density: a neural network, a Gaussian process, a random forest, a knowledge graph, a grammar, a decision process, a causal skeleton. Each is exposed through the same five-piece Distribution/Estimator/Accumulator/Sampler/Encoder contract (or, for the supervised/decision/causal ones, a small task-appropriate surface), so it composes with the stats core – a neural leaf drops into a CompositeDistribution, a GP into a mixture, and so on.

The right mental model is a small catalog of applied model families (mixle/models/README.md maps every module to one):

• neural & deep – neural nets, transformers/LMs, embeddings, and their training utilities;

• non-parametric – Gaussian processes and random forests as p(y | x) leaves;

• relational / structured – knowledge graphs, random graphs, grammars;

• latent-variable – Bayesian-nonparametric mixtures (Dirichlet process);

• decision & control – partially observable Markov decision processes;

• causal discovery – constraint-based structure learning.

(The imports below stay alphabetical – the ruff import sorter enforces that – so use the families above, not import order, as the map.)

These surfaces vary in maturity (see the Project status table in the top-level README); treat them as specialist adapters composable with the stable stats spine, not as the spine itself.

class LM(vocab, *, d_model=256, n_layer=6, n_head=8, block=128, device='cpu', embedding=None)

Bases: object

A causal-Transformer language model with a small declarative surface: fit / generate / nll.

Parameters:

• vocab (int)

• d_model (int)

• n_layer (int)

• n_head (int)

• block (int)

• device (str)

• embedding (Any)

to_dict()

Serialize the hyperparameters + trained weights so the LM survives a process boundary.

The token embedding may be tied across LMs (embedding=); from_dict rebuilds an untied module and loads the saved state_dict into it, so a round-tripped LM is standalone (any external tie is dropped).

Return type:

dict

classmethod from_dict(payload)

Rebuild an LM from to_dict() output (fresh module, saved weights loaded in).

Parameters:

payload (dict)

Return type:

LM

save(path)

Persist the LM’s inference artifact to path via torch.save: architecture config (vocab/d_model/n_layer/n_head/block) + learned weights, nothing else.

This is not a training checkpoint. It deliberately does not store the optimizer, LR scheduler, step count, RNG state, data-loader position, or gradient-scaler state, so load() returns a model ready for scoring and generation, not one you can resume training from without loss of state. For resumable training use the checkpoint path in mixle.utils.parallel.fault_tolerant_training / mixle.utils.parallel.dcp_checkpoint, which stores model + optimizer + data-loader position.

Parameters:

path (str)

Return type:

None

classmethod load(path)

Load an LM inference artifact written by save() (config + weights).

Returns a model ready for scoring/generation; this is not a training-resume checkpoint (see save()), so optimizer/step/RNG/data-loader state is not restored.

Loaded with weights_only=True: to_dict()’s payload is plain config scalars plus a tensor state_dict, never an arbitrary object graph, so this does not execute code from the file (unlike a full-module pickle – see mixle.models._neural_serial for that case).

Parameters:

path (str)

Return type:

LM

fit(token_ids, *, epochs=1, batch_size=64, lr=3e-3, distributed=False, precision='fp32', shuffle=True, dense=False, seed=0, log=None, tp_size=1, pp_size=1, cp_size=1, dp_replicate=1, dp_shard=1, ep_size=1, etp_size=1, microbatches=1, gradient_accumulation_steps=1, distributed_backend='torch_native', optimizer=None, max_grad_norm=None, compile=False, checkpoint_path=None, resume=False)

Small-scale reference pretraining (or continuation) on a token-id array.

Two paths, chosen by dense:

• dense=False (default, unchanged): the streaming estimator; the corpus is never buffered. It scores one next-token target per block-length window – the right shape for an unbounded stream, but only a fraction (~``1/block``) of the token supervision that packed dense teacher forcing extracts from the same tokens. A reference / continuation path for small-scale runs, not a frontier pretraining engine.

• dense=True: packed dense teacher forcing on a BOUNDED corpus. The token stream is cut into non-overlapping rows of block + 1 tokens (a partial tail row is dropped; there is no padding, so no attention-mask or EOS special-casing is needed – document boundaries are whatever the ids encode). Every row consumes block+1 tokens but contributes block shifted targets: the corpus’s first token and each cross-row transition are not supervised, and the partial tail is not consumed. last_fit_receipt reports those counts explicitly. Token ids stay integer end to end (no float conversion). seed drives row shuffling and log(epoch, mean_loss) reports per-epoch training loss, both as in fit_pairs(). With distributed=True this objective is executed by the selected distributed-gradient backend.

Parallel sizes are executable contracts. torch_native supports DDP/HSDP, FSDP2, MLP tensor parallelism, and CUDA context parallelism; unsupported combinations fail capability validation. Full TP/PP/CP/EP transformer and MoE training is exposed by the megatron backend.

Parameters:

• token_ids (Any)

• epochs (int)

• batch_size (int)

• lr (float)

• distributed (bool)

• precision (str)

• shuffle (bool)

• dense (bool)

• seed (int)

• log (Any)

• tp_size (int)

• pp_size (int)

• cp_size (int)

• dp_replicate (int)

• dp_shard (int)

• ep_size (int)

• etp_size (int)

• microbatches (int)

• gradient_accumulation_steps (int)

• distributed_backend (str)

• optimizer (Any)

• max_grad_norm (float | None)

• compile (bool)

• checkpoint_path (str | None)

• resume (bool)

Return type:

LM

fit_pairs(pairs, *, epochs=1, batch_size=32, lr=3e-3, mask_prompt=True, pad_id=0, seed=0, log=None)

Supervised fine-tuning on (prompt_ids, completion_ids) pairs with a dense per-position loss.

The streaming fit path scores ONE next-token target per window (the right shape for an unbounded pretraining stream); for a pair corpus that wastes a factor of block in compute. Here every position of every pair contributes cross-entropy in a single forward, and mask_prompt restricts the loss to completion positions – the standard SFT objective. Sequences longer than block keep the completion and drop the oldest prompt tokens; shorter ones are left-padded with pad_id (excluded from the loss). Include your end-of-sequence token in each completion so generate(stop_id=...) knows where to stop.

Parameters:

• pairs (Any)

• epochs (int)

• batch_size (int)

• lr (float)

• mask_prompt (bool)

• pad_id (int)

• seed (int)

• log (Any)

Return type:

LM

generate(prompt_ids, n=200, *, temperature=1.0, greedy=False, seed=0, stop_id=None)

Autoregressively extend prompt_ids by n tokens (greedy, or temperature-sampled).

stop_id ends generation early when that token is produced (it is included in the return value, so callers can strip it – and its presence distinguishes ‘finished’ from ‘ran out of budget’).

Parameters:

• prompt_ids (Any)

• n (int)

• temperature (float)

• greedy (bool)

• seed (int)

• stop_id (int | None)

Return type:

list

nll(token_ids)

Mean next-token negative log-likelihood (nats/token) on a token-id array.

Parameters:

token_ids (Any)

Return type:

float

class LMFitReceipt(objective, input_tokens, consumed_tokens_per_epoch, supervised_tokens_per_epoch, discarded_cross_row_transitions_per_epoch, dropped_tail_tokens, epochs)

Bases: object

Token-accounting receipt for one packed dense corpus layout.

Parameters:

• objective (str)

• input_tokens (int)

• consumed_tokens_per_epoch (int)

• supervised_tokens_per_epoch (int)

• discarded_cross_row_transitions_per_epoch (int)

• dropped_tail_tokens (int)

• epochs (int)

class CategoricalEmbedding(num_categories, dim, *, name=None, device='cpu', dtype='float32', init_seed=0)

Bases: object

A lazily-built learned embedding of shape (num_categories, dim); every consumer gets the same module.

Parameters:

• num_categories (int)

• dim (int)

• name (str | None)

• device (str)

• dtype (str)

• init_seed (int)

module()

The underlying nn.Embedding – built on first call, the identical instance thereafter.

Return type:

Any

class EmbeddingSpec(num_categories, dim, device, dtype, init_seed)

Bases: object

Immutable construction contract for a shared embedding table.

Parameters:

• num_categories (int)

• dim (int)

• device (str)

• dtype (str)

• init_seed (int)

class CausalSkeleton(edges, separating_sets, variable_names, all_separating_sets=None)

Bases: object

Undirected skeleton plus separating sets from a PC-style search.

Parameters:

• edges (set[tuple[int, int]])

• separating_sets (dict[tuple[int, int], frozenset[int]])

• variable_names (list[Any])

• all_separating_sets (dict[tuple[int, int], tuple[frozenset[int], ...]] | None)

has_edge(i, j)

Return whether the undirected skeleton contains edge i–j.

Parameters:

• i (int)

• j (int)

Return type:

bool

class CategoricalClassificationNeuralNetwork(module, engine=None, precision=None)

Bases: object

A Torch classifier wrapper optimized by summed categorical log likelihood.

The wrapped module must return one logits row per observation. Fitting is delegated to optimize_torch_objective so classification examples get the same convergence diagnostics and best-state restoration as distribution objectives.

Parameters:

• module (Any)

• engine (Any | None)

• precision (Any | None)

parameters()

Return trainable parameters of the wrapped classification module.

Return type:

Iterable[Any]

logits_tensor(x)

Return raw class logits for x as a Torch tensor.

Parameters:

x (Any)

Return type:

Any

log_likelihood(x, y)

Return the summed categorical log likelihood for integer labels.

Parameters:

• x (Any)

• y (Any)

Return type:

Any

fit(x, y, max_its=500, lr=0.01, optimizer='adam', tol=1.0e-7, out=None, print_iter=100, return_result=False, restore_best=True)

Maximize the categorical classification log likelihood.

Parameters:

• x (Any)

• y (Any)

• max_its (int)

• lr (float)

• optimizer (str)

• tol (float)

• out (Any | None)

• print_iter (int)

• return_result (bool)

• restore_best (bool)

Return type:

Any

predict_proba_tensor(x)

Return class probabilities for x as a Torch tensor.

Parameters:

x (Any)

Return type:

Any

predict_proba(x)

Return class probabilities for x as a NumPy array.

Parameters:

x (Any)

Return type:

ndarray

predict(x)

Return maximum-probability class labels for x.

Parameters:

x (Any)

Return type:

ndarray

class ConditionalIndependenceResult(measure, statistic, p_value, independent, status, sample_size, degrees_of_freedom, method, reason=None)

Bases: object

Result from a conditional independence calculation.

Parameters:

• measure (float | None)

• statistic (float | None)

• p_value (float | None)

• independent (bool | None)

• status (Literal['independent', 'dependent', 'inconclusive'])

• sample_size (int)

• degrees_of_freedom (int | None)

• method (str)

• reason (str | None)

discrete_conditional_independence(data, x, y, given=(), alpha=0.05)

Likelihood-ratio (G-test) conditional-independence test for categorical samples.

The statistic is 2 n I(X;Y|Z) and is compared with its chi-square null distribution. The asymptotic result is reported as inconclusive unless there are at least five observations per degree of freedom; callers that need sparse-table exactness should use a permutation test.

Parameters:

• data (Any)

• x (int)

• y (int)

• given (Sequence[int])

• alpha (float)

Return type:

ConditionalIndependenceResult

DPOLeaf

alias of DPOModel

class ErdosRenyiGraphModel(p, directed=False, self_loops=False, name=None)

Bases: object

Independent Bernoulli edge model for directed or undirected graphs.

Parameters:

• p (float)

• directed (bool)

• self_loops (bool)

• name (str | None)

classmethod fit_mle(adjacency, directed=False, self_loops=False, pseudo_count=0.0, prior_p=0.5, name=None)

Thin shim delegating to fit_erdos_renyi_mle (kept for the classmethod-fit call API).

Parameters:

• adjacency (Any)

• directed (bool)

• self_loops (bool)

• pseudo_count (float)

• prior_p (float)

• name (str | None)

Return type:

ErdosRenyiGraphModel

log_likelihood(adjacency)

Return the Bernoulli graph log likelihood.

Parameters:

adjacency (Any)

Return type:

float

sample(num_nodes, seed=None)

Draw one binary adjacency matrix.

Parameters:

• num_nodes (int)

• seed (int | None)

Return type:

ndarray

bic(adjacency)

Bayesian information criterion with one free parameter.

Parameters:

adjacency (Any)

Return type:

float

class GaussianProcessRegressor(lengthscale=1.0, amplitude=1.0, noise=0.1, mean=0.0, jitter=1.0e-6, kernel='rbf', engine=None, precision=None)

Bases: object

Exact GP regression with a stationary kernel and Gaussian observation noise.

The kernel is RBF (squared-exponential) by default; kernel="matern32" or "matern52" selects the Matern-3/2 or Matern-5/2 covariance, whose less smooth sample paths often fit physical responses better than the very smooth RBF.

Parameters:

• lengthscale (float)

• amplitude (float)

• noise (float)

• mean (float)

• jitter (float)

• kernel (str)

• engine (Any | None)

• precision (Any | None)

parameters()

Return trainable raw kernel/noise parameters and the mean.

property lengthscale: float

Return the fitted kernel lengthscale.

property amplitude: float

Return the fitted kernel amplitude.

property noise: float

Return the fitted Gaussian observation-noise standard deviation.

kernel(x1, x2)

Return the covariance matrix between two input arrays under the configured kernel.

Parameters:

• x1 (Any)

• x2 (Any)

Return type:

Any

log_marginal_likelihood(x, y)

Return the exact GP log marginal likelihood for training data.

Parameters:

• x (Any)

• y (Any)

Return type:

Any

fit(x, y, max_its=500, lr=0.05, optimizer='adam', tol=1.0e-7, out=None, print_iter=100, return_result=False, restore_best=True)

Maximize the GP log marginal likelihood and return self.

Returns the fitted model so model = gp.fit(x, y) works like every other Model.fit in mixle.models. Set return_result=True for the full objective diagnostics (an ObjectiveFitResult carrying the objective value, iteration count, and history).

Compatibility note: before 0.8.0 the default return was the (value, iterations) tuple; those live behind return_result=True now (result.value / result.iterations).

Parameters:

• x (Any)

• y (Any)

• max_its (int)

• lr (float)

• optimizer (str)

• tol (float)

• out (Any | None)

• print_iter (int)

• return_result (bool)

• restore_best (bool)

Return type:

Any

predict(x_train, y_train, x_new, return_cov=False)

Return posterior predictive mean, and optionally covariance.

Parameters:

• x_train (Any)

• y_train (Any)

• x_new (Any)

• return_cov (bool)

Return type:

Any

predict_monotone(x_train, y_train, x_new, increasing=True)

Return the posterior-mean prediction projected to be monotone in scalar x_new.

Predicts the GP posterior mean at x_new and projects it onto the monotone cone (non-decreasing if increasing else non-increasing) by pool-adjacent-violators in x_new order – the L2-closest monotone curve to the GP mean. Intended for scalar (1-D) inputs (e.g. monotone age-depth / dose-response fits); reduces to predict() when the posterior mean is already monotone.

Parameters:

• x_train (Any)

• y_train (Any)

• x_new (Any)

• increasing (bool)

Return type:

ndarray

class SparseGaussianProcessRegressor(lengthscale=1.0, amplitude=1.0, noise=0.1, kernel='rbf', n_inducing=50, jitter=1e-8)

Bases: object

Sparse GP regression with m inducing points (FITC).

Parameters:

• lengthscale – initial kernel/noise hyperparameters (all positive and finite).

• amplitude – initial kernel/noise hyperparameters (all positive and finite).

• noise – initial kernel/noise hyperparameters (all positive and finite).

• kernel – 'rbf', 'matern32' or 'matern52'.

• n_inducing – number of inducing points (placed by k-means++ over unique training inputs at fit).

• jitter – strictly positive diagonal stabilization for the inducing covariance.

fit(x, y, *, optimize=True, seed=0, max_iter=100)

Place inducing points and (optionally) fit hyperparameters by the FITC marginal likelihood.

predict(x_new, *, return_var=False)

Posterior mean (and optionally marginal variance) at x_new. O(m^2) per query batch.

class GaussianRegressionNeuralNetwork(module, noise=1.0, engine=None, precision=None)

Bases: object

A Torch module trained with a Gaussian regression log likelihood.

The wrapped module predicts the response mean and this helper learns a scalar observation noise alongside module weights. It uses the same generic Torch objective optimizer as the distribution objective helpers.

Parameters:

• module (Any)

• noise (float)

• engine (Any | None)

• precision (Any | None)

parameters()

Return trainable module parameters plus the raw noise parameter.

Return type:

Iterable[Any]

property noise: float

Return the fitted observation standard deviation.

predict_tensor(x)

Return module predictions as a Torch tensor on the configured engine.

Parameters:

x (Any)

Return type:

Any

log_likelihood(x, y)

Return the summed Gaussian regression log likelihood.

Parameters:

• x (Any)

• y (Any)

Return type:

Any

fit(x, y, max_its=500, lr=0.01, optimizer='adam', tol=1.0e-7, out=None, print_iter=100, return_result=False, restore_best=True)

Maximize the Gaussian regression log likelihood.

The default return shape is the historical (value, iterations) tuple. Set return_result=True for the full objective diagnostics.

Parameters:

• x (Any)

• y (Any)

• max_its (int)

• lr (float)

• optimizer (str)

• tol (float)

• out (Any | None)

• print_iter (int)

• return_result (bool)

• restore_best (bool)

Return type:

Any

predict(x)

Return mean predictions as a NumPy array.

Parameters:

x (Any)

Return type:

ndarray

class NeuralCategorical(module, m_steps=40, lr=0.01, name=None, batch_size=None, device='cpu', max_optimizer_steps=None, optimizer_state=None)

Bases: SequenceEncodableProbabilityDistribution

p(y | x) = softmax(module(x)) as a mixle leaf. Observation is the pair (x, y), y an int class.

batch_size (None = full batch) makes the M-step minibatch SGD over m_steps passes – needed to train a real conv net on a large image set; max_optimizer_steps optionally caps updates independently of batch size, and device (e.g. "mps"/"cuda") runs them on the GPU.

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• name (str | None)

• batch_size (int | None)

• device (str)

• max_optimizer_steps (int | None)

• optimizer_state (dict[str, Any] | None)

log_density(xy)

Return log p(y | x) for one feature/class observation pair.

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)

Return per-row categorical conditional log probabilities for encoded pairs.

Parameters:

enc (Any)

Return type:

ndarray

predict(x)

Return maximum-probability class predictions for one or more inputs.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a conditional sampler over labels given features.

Parameters:

seed (int | None)

Return type:

NeuralCategoricalSampler

estimator(pseudo_count=None)

Return the generalized-EM estimator for weighted cross-entropy training.

Parameters:

pseudo_count (float | None)

Return type:

NeuralCategoricalEstimator

dist_to_encoder()

Return the encoder for (x, class) observation pairs.

Return type:

NeuralCategoricalEncoder

to_dict()

Serialize hyperparameters and module bytes for registry-based round trips.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a NeuralCategorical from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralCategorical

class NeuralConditionalDensity(module, *, m_steps=60, lr=5e-3, device='cpu', name=None)

Bases: SequenceEncodableProbabilityDistribution

Wrap a torch conditional-density module (module.log_density(x, y) -> (n,)) as a mixle leaf.

Observations are pairs (x, y). The module must also expose sample_given(x) -> (n, d) to draw y.

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• name (str | None)

log_density(xy)

Return log p(y | x) for one observation pair (x, y).

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)

Return per-row conditional log densities for encoded (x, y) arrays.

Parameters:

enc (Any)

Return type:

ndarray

sampler(seed=None)

Return a conditional sampler for drawing y given x.

Parameters:

seed (int | None)

Return type:

NeuralConditionalDensitySampler

estimator(pseudo_count=None)

Return the generalized-EM estimator for weighted conditional-density training.

Parameters:

pseudo_count (float | None)

Return type:

NeuralConditionalDensityEstimator

dist_to_encoder()

Return the encoder for (x, y) observation pairs.

Return type:

NeuralConditionalDensityEncoder

to_dict()

Serialize hyperparameters and module bytes for registry-based round trips.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a NeuralConditionalDensity from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralConditionalDensity

class GradEstimator(module, *, m_steps=60, lr=5e-3, device=None, batch_size=None, max_optimizer_steps=None, precision='auto', name=None, loss=None, optimizer=None, lr_decay=None)

Bases: ParameterEstimator

M-step: responsibility-weighted MLE – max sum_i w_i log p(x_i) by gradient ascent on the module (warm-started across EM iterations). loss/optimizer are the caller’s hooks; the optimizer only ever sees trainable parameters, so frozen submodules stay frozen and a fully frozen module makes the M-step a no-op (a fixed distribution).

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (Any)

• batch_size (int | None)

• max_optimizer_steps (int | None)

• precision (str)

• name (str | None)

• loss (Any)

• optimizer (Any)

• lr_decay (float | None)

accumulator_factory()

Return the accumulator factory used to collect this estimator’s sufficient statistics.

Return type:

DataBufferAccumulatorFactory

estimate(nobs, suff_stat)

Estimate a distribution from accumulated sufficient statistics.

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

GradLeaf

class GradLeaf(module, *, m_steps=60, lr=5e-3, device=None, batch_size=None, max_optimizer_steps=None, precision='auto', name=None, loss=None, optimizer=None, lr_decay=None)

Bases: SequenceEncodableProbabilityDistribution

Wrap a torch density module (module.log_density(x) -> (n,)) as a composable mixle

Measured negative, so nobody re-derives it (2026-07-12, Apple M4, torch 2.12 CPU): wrapping the full-batch M-step loss in torch.compile is a LOSS here – 0.93x at a 2x32 MLP / n=100k and 0.79x at 2x256 / n=200k, plus ~6s compile overhead per module – so there is deliberately no compile= flag. Re-measure before adding one (a CUDA build or a much larger module could flip it); the probe script pattern lives in the introducing PR. distribution (see the module docstring). loss and optimizer are the M-step hooks.

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (Any)

• batch_size (int | None)

• max_optimizer_steps (int | None)

• precision (str)

• name (str | None)

• loss (Any)

• optimizer (Any)

• lr_decay (float | None)

log_density(x)

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

Parameters:

x (Any)

Return type:

float

seq_log_density(x)

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

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

GradLeafSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

GradEstimator

dist_to_encoder()

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

Return type:

GradLeafEncoder

class NeuralDensity(module, *, m_steps=60, lr=5e-3, device='cpu', precision='auto', name=None)

Bases: GradLeaf

Wrap a torch density module (module.log_density(x) -> (n,)) as a composable mixle distribution.

A thin named subclass of GradLeaf – the generic bridge owns the manufactured contract (buffer accumulator, array encoder, gradient M-step, sampler); this class owns only its name, its JSON payload, and its ready-module builders below. loss/optimizer hooks pass through (see the grad_leaf module docstring for the control story).

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• precision (str)

• name (str | None)

log_density(x)

Return log p(x) for one observation under the wrapped density module.

Parameters:

x (Any)

Return type:

float

seq_log_density(x)

Return per-row log densities for encoded observations.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a sampler delegating to the wrapped module’s sample method.

Parameters:

seed (int | None)

Return type:

NeuralDensitySampler

estimator(pseudo_count=None)

Return the generalized-EM estimator for weighted neural-density training.

Parameters:

pseudo_count (float | None)

Return type:

NeuralDensityEstimator

dist_to_encoder()

Return the encoder for vectorized neural-density scoring and fitting.

Return type:

NeuralDensityEncoder

to_dict()

Serialize hyperparameters and module bytes for registry-based round trips.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a NeuralDensity from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralDensity

class NeuralDensityEstimator(module, *, m_steps=60, lr=5e-3, device='cpu', precision='auto', name=None)

Bases: GradEstimator

M-step: responsibility-weighted MLE – max sum_i w_i log p(x_i) by gradient ascent on the module (warm).

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• precision (str)

• name (str | None)

accumulator_factory()

Return an accumulator factory for weighted neural-density batches.

Return type:

DataBufferAccumulatorFactory

class NeuralGaussian(module, noise=1.0, m_steps=40, lr=0.01, name=None, device=None)

Bases: SequenceEncodableProbabilityDistribution

p(y | x) = N(y; module(x), noise^2 I) as a mixle leaf. Observation is the pair (x, y).

Parameters:

• module (Any)

• noise (float)

• m_steps (int)

• lr (float)

• name (str | None)

• device (Any)

log_density(xy)

Return log p(y | x) for one encoded observation pair (x, y).

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)

Return per-row Gaussian conditional log densities for encoded (x, y) arrays.

Parameters:

enc (Any)

Return type:

ndarray

classmethod compute_capabilities()

Declare engine-ready scoring support for NumPy and Torch execution backends.

backend_seq_log_density(enc, engine)

Engine-neutral vectorized log-density for encoded (x, y) pairs.

Parameters:

• enc (Any)

• engine (Any)

Return type:

Any

sampler(seed=None)

Return a conditional sampler for drawing y given x.

Parameters:

seed (int | None)

Return type:

NeuralGaussianSampler

estimator(pseudo_count=None)

Return the generalized-EM estimator for responsibility-weighted neural regression.

Parameters:

pseudo_count (float | None)

Return type:

NeuralGaussianEstimator

dist_to_encoder()

Return the encoder for batches of (x, y) observation pairs.

Return type:

NeuralGaussianEncoder

to_dict()

Serialize hyperparameters and module bytes for registry-based round trips.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a NeuralGaussian from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralGaussian

class VAE(dim, *, latent=2, hidden=32, eval_samples=16, m_steps=120, lr=5e-3, device='cpu', name=None)

Bases: _NeuralFamily

A latent-variable p(x) over R^dim via a variational autoencoder.

log_density is a reproducible multi-sample estimate of the ELBO, with uncertainty available from the wrapped module’s elbo_estimate method. It is not an exact log density. See build_vae() for the full statement.

Parameters:

• dim (int)

• latent (int)

• hidden (int)

• eval_samples (int)

• m_steps (int)

• lr (float)

• device (str)

• name (str | None)

class Flow(dim, *, hidden=32, layers=4, m_steps=80, lr=5e-3, device='cpu', name=None)

Bases: _NeuralFamily

An exact p(x) over R^dim via a RealNVP coupling flow (invertible map to a standard-normal base).

Parameters:

• dim (int)

• hidden (int)

• layers (int)

• m_steps (int)

• lr (float)

• device (str)

• name (str | None)

class MAF(dim, *, hidden=64, blocks=3, m_steps=80, lr=5e-3, device='cpu', name=None)

Bases: _NeuralFamily

An exact p(x) over R^dim via a masked autoregressive flow (richer autoregressive dependence).

Parameters:

• dim (int)

• hidden (int)

• blocks (int)

• m_steps (int)

• lr (float)

• device (str)

• name (str | None)

class DiscreteAR(dim, cats, *, hidden=64, m_steps=100, lr=5e-3, device='cpu', name=None)

Bases: _NeuralFamily

An exact, normalized p(x) over discrete vectors x in {0..cats-1}^dim (autoregressive, MADE-masked).

Parameters:

• dim (int)

• cats (int)

• hidden (int)

• m_steps (int)

• lr (float)

• device (str)

• name (str | None)

class DPOModel(policy, ref, beta=0.1, m_steps=100, lr=1e-3, device='cpu')

Bases: SequenceEncodableProbabilityDistribution

DPO over (x, chosen, rejected) preference triples. policy is trained, ref is frozen.

Parameters:

• policy (Any)

• ref (Any)

• beta (float)

• m_steps (int)

• lr (float)

• device (str)

seq_log_density(enc)

Return per-row DPO preference log likelihoods for encoded triples.

Parameters:

enc (Any)

Return type:

ndarray

log_density(xcr)

Return the DPO log likelihood for one (x, chosen, rejected) triple.

Parameters:

xcr (Any)

Return type:

float

prefers(x)

The policy’s argmax action at x – what the aligned policy now picks.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return the sampler for the preference-scoring leaf.

Parameters:

seed (int | None)

Return type:

DPOModelSampler

estimator(pseudo_count=None)

Return the DPO estimator that trains the policy while keeping the reference fixed.

Parameters:

pseudo_count (float | None)

Return type:

DPOModelEstimator

dist_to_encoder()

Return the encoder for preference triples.

Return type:

DPOEncoder

to_dict()

Serialize policy/reference modules and DPO hyperparameters.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a DPOModel from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

DPOModel

class EnergyModel(module, *, m_steps=200, lr=5e-3, noise_ratio=1, langevin_steps=40, langevin_step=0.05, device='cpu', name=None)

Bases: SequenceEncodableProbabilityDistribution

log p(x) ≈ -E(x) + c for an energy module.

The module exposes energy(x) -> (n,) and a learned scalar log_norm whose convention is c=-log Z. Approximately normalized (trained by NCE); log_density returns -E(x) + c. Composes like any leaf.

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• noise_ratio (int)

• langevin_steps (int)

• langevin_step (float)

• device (str)

• name (str | None)

log_density(x)

Return the approximate normalized log density for one observation.

Parameters:

x (Any)

Return type:

float

seq_log_density(x)

Return approximate normalized log densities for a batch of observations.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a Langevin sampler for the learned energy model.

Parameters:

seed (int | None)

Return type:

EnergyModelSampler

estimator(pseudo_count=None)

Return the NCE estimator used as the model’s M-step.

Parameters:

pseudo_count (float | None)

Return type:

EnergyModelEstimator

dist_to_encoder()

Return the encoder for vectorized energy-model scoring.

Return type:

EnergyModelEncoder

to_dict()

Serialize hyperparameters and module bytes for registry-based round trips.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild an EnergyModel from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

EnergyModel

class CoarseningMetrics(source_parameter_count, result_parameter_count, source_parameter_bytes, result_parameter_bytes, source_manifest_entries, result_manifest_entries, source_leaf_blocks, result_leaf_blocks, source_estimated_block_evaluations, result_estimated_block_evaluations, latency_seconds=None, latency_basis='unmeasured: no calibration input was supplied')

Bases: object

Structural measurements for the source and returned executable models.

estimated_block_evaluations counts ordinary block-branch evaluations in one forward: a plain block costs one and the current numerical Taylor merge costs four. FLOPs and latency are deliberately not fabricated without shapes or calibration inputs.

Parameters:

• source_parameter_count (int)

• result_parameter_count (int)

• source_parameter_bytes (int)

• result_parameter_bytes (int)

• source_manifest_entries (int)

• result_manifest_entries (int)

• source_leaf_blocks (int)

• result_leaf_blocks (int)

• source_estimated_block_evaluations (int)

• result_estimated_block_evaluations (int)

• latency_seconds (float | None)

• latency_basis (str)

class CoarsenResult(model, receipt_map=<factory>, accepted_pairs=<factory>, rejected_pairs=<factory>, depth_only_kl=0.0, total_kl=0.0, budget=inf, trust_region=inf, within_budget=True, structure_receipts=<factory>, artifact_digest=None, metrics=None)

Bases: object

Output of coarsen(): the new (shallower) model, the full per-scale receipt map, and the bookkeeping needed to see exactly which merges were accepted vs. rejected and why.

structure_receipts is the third move’s own receipt list – see _narrow_block_mlp2() – kept OUT of receipt_map deliberately: receipt_map values are ScaleReceipt (closed-form KL against a Gaussian law, consumed as-is by hybrid’s surrogate_closure_error-keyed stage ranking in mixle.models.compress), while structure- projection’s own receipt is a ProjectionReceipt (a Sigma-weighted reconstruction error, not a KL) – mixing the two dataclasses into one dict would silently break that attribute lookup. The final_artifact scale receipt is authoritative for the returned model: its end-to-end law, total_kl, and digest are computed after every accepted structure projection.

Parameters:

• model (Any)

• receipt_map (dict[str, ScaleReceipt])

• accepted_pairs (list[tuple[int, int]])

• rejected_pairs (list[tuple[int, int]])

• depth_only_kl (float)

• total_kl (float)

• budget (float)

• trust_region (float)

• within_budget (bool)

• structure_receipts (list[ProjectionReceipt])

• artifact_digest (str | None)

• metrics (CoarseningMetrics | None)

class ProjectionReceipt(name, mode, sigma_weighted_error, accepted=True, final_kl_divergence=None, artifact_digest=None)

Bases: object

Receipt for a structure_project() call – reports G2’s own Sigma-weighted reconstruction error directly (there is no Gaussian law on either side of a weight-space projection, so this is not a KL divergence; it is the SAME sigma_weighted_error objective G2’s solvers themselves minimize).

Parameters:

• name (str)

• mode (str)

• sigma_weighted_error (float)

• accepted (bool)

• final_kl_divergence (float | None)

• artifact_digest (str | None)

class ScaleReceipt(name, teacher_law, student_law, kl_divergence, surrogate_closure_error, accepted=True, divergence_basis='full_rank_gaussian', regularized_kl_divergence=None, assumptions=(), sensitivity=<factory>, scope='local_candidate', artifact_digest=None)

Bases: object

One per-scale receipt: the CLOSED-FORM teacher/student divergence at this coarsening step, plus (separately) G1’s own closure-error signal for how much the Gaussian-surrogate assumption itself is trusted at this point in the network (nan where no G1 block closure applies, e.g. width-merge, which never runs a real Block forward and so has nothing for G1’s Monte-Carlo closure check to compare against).

Parameters:

• name (str)

• teacher_law (MultivariateGaussianDistribution)

• student_law (MultivariateGaussianDistribution)

• kl_divergence (float)

• surrogate_closure_error (float)

• accepted (bool)

• divergence_basis (str)

• regularized_kl_divergence (float | None)

• assumptions (tuple[str, ...])

• sensitivity (dict[str, float])

• scope (str)

• artifact_digest (str | None)

class CompressedModel(model, method, receipt, non_sampling_receipts=<factory>, hybrid_selected_stages=<factory>, hybrid_sample_indices=<factory>)

Bases: object

Unified result: the compressed module plus the receipt explaining which method produced it, and (when the chosen/underlying path touched non_sampling) the raw per-stage ScaleReceipt map that path’s closure-error signal came from.

Parameters:

• model (Any)

• method (str)

• receipt (CompressionReceipt)

• non_sampling_receipts (dict[str, ScaleReceipt])

• hybrid_selected_stages (list[str])

• hybrid_sample_indices (list[int])

class CompressionReceipt(method, auto, quality, sample_count, selection_training_sample_count=0, selection_evaluation_sample_count=0, quality_basis='unmeasured', candidates=<factory>, notes='')

Bases: object

Explains why method was used: its own measured numbers, plus – for auto-pick – the same real numbers for every OTHER method considered and rejected. Mirrors mixle.models.unified_quantizer.QuantizationReceipt.

Parameters:

• method (str)

• auto (bool)

• quality (float | None)

• sample_count (int)

• selection_training_sample_count (int)

• selection_evaluation_sample_count (int)

• quality_basis (str)

• candidates (dict[str, MethodCandidate])

• notes (str)

property total_selection_sample_count: int

All real training plus evaluation examples consumed to make an automatic choice.

class MethodCandidate(method, quality, sample_count, evaluation_sample_count, reward)

Bases: object

Real, measured numbers for one method family considered for one compress() call – the raw material every CompressionReceipt (chosen or rejected) is built from, mirroring mixle.models.unified_quantizer.MethodCandidate.

Parameters:

• method (str)

• quality (float)

• sample_count (int)

• evaluation_sample_count (int)

• reward (float)

class EvalReport(checkpoint_id, tasks, seed, metadata=<factory>)

Bases: object

The per-rung receipt: every task’s score for one checkpoint, JSON-serializable via report().

Parameters:

• checkpoint_id (str)

• tasks (list[TaskResult])

• seed (int)

• metadata (dict[str, Any])

scores()

{task_name: score} – the compact view track_regression() consumes.

Return type:

dict[str, float | None]

class RegressionFlag(task, checkpoint_id, checkpoint_index, current_score, reference_score, reference_checkpoint_id, reference_index, relative_delta, threshold)

Bases: object

One metric that regressed beyond threshold relative to its best-so-far value in the sequence.

Parameters:

• task (str)

• checkpoint_id (str)

• checkpoint_index (int)

• current_score (float)

• reference_score (float)

• reference_checkpoint_id (str)

• reference_index (int)

• relative_delta (float)

• threshold (float)

class RegressionReport(flags, n_checkpoints, threshold, reference)

Bases: object

The regression-tracking receipt over an ordered sequence of EvalReport.

Parameters:

• flags (list[RegressionFlag])

• n_checkpoints (int)

• threshold (float)

• reference (str)

class TaskResult(name, score, higher_is_better, n_examples, details=<factory>, succeeded=True, error=None)

Bases: object

One capability-axis score: a name, a scalar, and which direction is “better” for regression math.

Parameters:

• name (str)

• score (float | None)

• higher_is_better (bool)

• n_examples (int)

• details (dict[str, Any])

• succeeded (bool)

• error (str | None)

class CompressedAdam(*_args, **_kwargs)

Bases: object

Placeholder for the torch-only optimizer; constructing it reports the missing dependency.

Parameters:

• _args (Any)

• _kwargs (Any)

class CompressedMomentEncoding(method, shape, g4_encoding=None, int8_codes=None, int8_scales=None, int8_block_size=2048, dense_values=None, source_dtype='<f4', schema_version=1, requested_method='auto', fallback_reason=None)

Bases: object

Storage format for ONE moment tensor (m or v), compressed by exactly one of the three available methods – only the fields for method are populated, mirroring SortedProfileEncoding’s single-active-branch convention.

Parameters:

• method (str)

• shape (tuple)

• g4_encoding (SortedProfileEncoding | None)

• int8_codes (ndarray | None)

• int8_scales (ndarray | None)

• int8_block_size (int)

• dense_values (ndarray | None)

• source_dtype (str)

• schema_version (int)

• requested_method (str)

• fallback_reason (str | None)

method

"g4" (sorted-profile, mixle.models.sorted_profile_quantizer), "int8" (blockwise quantization), or "dense" (fallback – neither compressed representation was trustworthy for this tensor).

Type:

str

shape

Original tensor shape.

Type:

tuple[int, …]

g4_encoding

Populated iff method == "g4".

Type:

mixle.models.sorted_profile_quantizer.SortedProfileEncoding | None

int8_codes / int8_scales / int8_block_size

Populated iff method == "int8".

dense_values

Populated iff method == "dense".

Type:

numpy.ndarray | None

nbytes()

Measured storage footprint, in bytes – delegates to G4’s own receipt-carrying nbytes() for the G4 branch; computes int8/dense directly.

Return type:

int

to_tensor_state()

Return a portable tensor tree with no custom Python objects or NumPy arrays.

Return type:

dict[str, Any]

classmethod from_tensor_state(state)

Rebuild and validate an encoding from to_tensor_state() output.

Parameters:

state (Any)

Return type:

CompressedMomentEncoding

class CompressedOptimizerState(shape, method='auto', *, state_dtype=None, **compress_kwargs)

Bases: object

Adam-style m/v moment storage for ONE parameter tensor, held COMPRESSED between optimizer steps (rather than as two dense fp32 buffers).

set compresses; get decompresses back to plain tensors for the optimizer math to use.

v (Adam’s second moment, an EMA of squared gradients) is stored as its COMPRESSED SQUARE ROOT rather than compressed directly. This is a real, load-bearing design choice, not cosmetic: v routinely spans many orders of magnitude within one parameter tensor (a few large-gradient elements next to many near-zero ones), which is exactly the dynamic range that defeats both compression paths – linear int8 quantization sets one block-wide scale from the block’s max, so any element more than ~127x smaller than that max quantizes to LITERALLY zero; G4’s Gaussian tail fit is symmetric and can reconstruct small true values as slightly negative. Either failure, fed straight into Adam’s 1/(sqrt(v_hat) + eps) denominator, produces a step blown up by orders of magnitude (empirically confirmed: an early version of this module, compressing v directly with int8, diverged to a >100x loss spike within ~5 steps on the tiny transformer this module’s own test suite trains). sqrt(v) roughly halves the dynamic range in log-space (Adam’s own second-moment buffer already tracks squared gradients FOR this reason – sqrt(v) is the RMS gradient-magnitude scale, the quantity Adam’s update actually normalizes by), and squaring the reconstructed value back on get() is a nonnegative projection for free – it can never reconstruct a negative v, closing the negative-v/NaN failure mode without a separate clamp needing to paper over it (a defensive clamp_min(0.0) is still applied as a last line of defense in CompressedAdam, since callers of CompressedOptimizerState directly should not have to know this internal detail to stay safe).

Parameters:

• shape (tuple)

• method (str)

• state_dtype (Any)

• compress_kwargs (Any)

nbytes()

Total measured compressed footprint of both moment buffers, in bytes.

Return type:

int

property methods: tuple[str | None, str | None]

(m_method, v_method) – the ACTUAL method used for each buffer (post any honest downgrade-to-dense), not just the one requested.

to_tensor_state()

Return a resumable tensor tree suitable for safe optimizer checkpoints.

Return type:

dict[str, Any]

class Fp8CastResult(tensor, used_fp8, reason, max_abs, underflow_fraction)

Bases: object

Receipt of one fp8_cast_with_guard() call – never silently swallowed.

Parameters:

• tensor (Any)

• used_fp8 (bool)

• reason (str)

• max_abs (float)

• underflow_fraction (float)

tensor

The output tensor – either the fp8-cast tensor (used_fp8=True) or the fallback_dtype cast (used_fp8=False).

Type:

Any

used_fp8

Whether the fp8 cast was accepted.

Type:

bool

reason

Human-readable reason for the decision (acceptance or the specific guard that fired).

Type:

str

max_abs

The input tensor’s max absolute value (0.0 for an empty tensor) – the statistic the overflow guard checks.

Type:

float

underflow_fraction

Fraction of the input’s nonzero values that flushed to exactly zero under the fp8 round-trip – the statistic the underflow guard checks. 0.0 when the overflow guard fired first (the round-trip was never attempted).

Type:

float

class RecomputeDecision(block_index, should_recompute, estimated_cost, estimated_benefit, activation_bytes, recompute_flops, rationale)

Bases: object

The cost/benefit tradeoff and chosen action for ONE block’s activation-recompute policy – mirrors D6’s RespecializationDecision shape: a flat, inspectable dataclass carrying both the raw estimates and the derived decision, not just a boolean.

Parameters:

• block_index (int)

• should_recompute (bool)

• estimated_cost (float)

• estimated_benefit (float)

• activation_bytes (float)

• recompute_flops (float)

• rationale (str)

property net_benefit: float

estimated_benefit - estimated_cost – positive iff recomputing this block is worth it.

class SelectiveRecomputePolicy(memory_value_per_byte=_DEFAULT_MEMORY_VALUE_PER_BYTE, flop_cost_per_unit=_DEFAULT_FLOP_COST_PER_UNIT)

Bases: object

Per-block, cost-model-driven activation-checkpointing decision – extends mixle.models.transformer.CausalLM’s previously all-or-nothing gradient_checkpointing bool flag (see that module’s forward, which now also accepts a per-block list) to a PER-BLOCK decision.

A block is recommended for recompute when the value of the memory freed (its stored- activation footprint, valued at memory_value_per_byte) exceeds the cost of the extra compute spent recomputing it (its recompute FLOPs, valued at flop_cost_per_unit) – the same cost-vs-benefit tradeoff SHAPE as D6’s compile economics, applied to this different decision (recompute-vs-store, not eager-vs-compiled).

Parameters:

• memory_value_per_byte (float)

• flop_cost_per_unit (float)

decide_model(lm, batch, seq_len, dtype_bytes=4)

Decide per-block recompute for every block of a CausalLM (mixle.models.transformer.build_causal_lm), bound to its actual block manifest.

Linear recompute cost is measured from every block’s real parameter count, so edited, widened, coarsened, or otherwise heterogeneous blocks can receive different decisions.

Parameters:

• lm (Any)

• batch (int)

• seq_len (int)

• dtype_bytes (int)

Return type:

list[RecomputeDecision]

apply_to_model(lm, batch, seq_len, dtype_bytes=4)

Compute the per-block decisions and set them directly as lm.gradient_checkpointing (a per-block bool list – see mixle.models.transformer.CausalLM.forward, which accepts either a single bool for all blocks or a per-block list).

Parameters:

• lm (Any)

• batch (int)

• seq_len (int)

• dtype_bytes (int)

Return type:

list[RecomputeDecision]

GaussianLaw

alias of MultivariateGaussianDistribution

class GeluCovarianceReceipt(correction_frobenius, min_eigenvalue_before, min_eigenvalue_after, iterations, marginal_variance_max_error)

Bases: object

Measured correction applied to keep linearized GELU correlations PSD at exact marginal variances.

Parameters:

• correction_frobenius (float)

• min_eigenvalue_before (float)

• min_eigenvalue_after (float)

• iterations (int)

• marginal_variance_max_error (float)

class LayerMoments(index, name, law, closure_error, diagnostics=<factory>)

Bases: object

One layer’s propagated law plus its locally-computed closure-error receipt.

Parameters:

• index (int)

• name (str)

• law (MultivariateGaussianDistribution)

• closure_error (float)

• diagnostics (dict[str, Any])

class NeuralOptimizerPlan(routes, sign_stable=None)

Bases: object

Inspectable non-Adam-first routes for a Torch module.

Parameters:

• routes (tuple[NeuralOptimizerRoute, ...])

• sign_stable (bool | None)

class NeuralOptimizerRoute(name, family, reason, shape)

Bases: object

One named parameter’s executable update family and selection reason.

Parameters:

• name (str)

• family (str)

• reason (str)

• shape (tuple[int, ...])

class RoutedNeuralOptimizer(groups, defaults, plan, *, precondition_frequency, muon_backend, muon_steps, fisher_provider, projection_provider)

Bases: object

Checkpointable mixed-family optimizer with one validated route per trainable parameter.

Parameters:

• groups (list[dict[str, Any]])

• defaults (dict[str, Any])

• plan (Any)

• precondition_frequency (int)

• muon_backend (str)

• muon_steps (int)

• fisher_provider (Callable[[str, Any], Any] | None)

• projection_provider (Callable[[str, Any], None] | None)

class EMATeacher(model, decay=0.999)

Bases: object

An exponential-moving-average copy of a model’s own weights: the standard mean-teacher/BYOL/DINO self-distillation teacher.

update(student_model) applies teacher = decay * teacher + (1 - decay) * student to every tensor in the teacher’s state_dict (parameters AND buffers, so e.g. non-trainable statistics stay consistent too) – called once per training step, AFTER the optimizer step, so the teacher always tracks a temporally-smoothed trailing average of the student. The teacher is a real, independent, forward-passable module (forward/predict), held in eval mode with gradients disabled: it is a read-only distillation TARGET, never itself directly optimized.

Parameters:

• model (Any)

• decay (float)

update(student_model)

One EMA step: pull every teacher tensor toward the student’s current value.

CausalLM ties head.weight to tok.weight (weight tying), so several state_dict() keys alias the SAME underlying storage – updating each key naively would apply the EMA formula to that storage more than once per step (a double update). seen dedupes by storage identity (data_ptr()) so every real tensor is updated exactly once, however many names alias it.

Parameters:

student_model (Any)

Return type:

None

forward(x)

Run x through the EMA-teacher weights (eval mode, no grad).

Parameters:

x (Any)

Return type:

Any

predict(x)

Run x through the EMA-teacher weights (eval mode, no grad).

Parameters:

x (Any)

Return type:

Any

class TrainStats(ce_loss=<factory>, stochastic_depth_loss=<factory>, ema_consistency_loss=<factory>, total_loss=<factory>)

Bases: object

Per-step telemetry from train_with_self_distillation() – cross-entropy, stochastic-depth consistency, and EMA-teacher consistency losses, kept separately so a caller can see which pressure is doing what (and the combined total actually optimized).

Parameters:

• ce_loss (list[float])

• stochastic_depth_loss (list[float])

• ema_consistency_loss (list[float])

• total_loss (list[float])

class FeatureMapDensity(feature_name, inner, name=None, *, feature_version=None, feature_digest=None, feature_dim=None)

Bases: SequenceEncodableProbabilityDistribution

p(feature_fn(x)) for a registered, deterministic feature_fn and inner distribution.

Parameters:

• feature_name (str)

• inner (Any)

• name (str | None)

• feature_version (str | None)

• feature_digest (str | None)

• feature_dim (int | None)

density(x)

Return the induced feature-space density at raw item x.

Parameters:

x (Any)

Return type:

float

log_density(x)

Return log p(feature_fn(x)) under the inner distribution.

Parameters:

x (Any)

Return type:

float

seq_log_density(x)

Return inner log densities for an already-featurized batch.

Parameters:

x (ndarray)

Return type:

ndarray

sampler(seed=None)

Raw-item sampling is undefined because the feature map has no inverse.

Parameters:

seed (int | None)

Return type:

DistributionSampler

feature_sampler(seed=None)

Return the sampler of the explicitly separate inner feature-space distribution.

Parameters:

seed (int | None)

Return type:

DistributionSampler

estimator(pseudo_count=None)

Return an estimator that fits the inner estimator on registered features.

Parameters:

pseudo_count (float | None)

Return type:

FeatureMapEstimator

dist_to_encoder()

Return the encoder that maps raw items to feature vectors.

Return type:

FeatureMapEncoder

class FeatureMapEstimator(feature_name, inner, name=None, *, feature_version=None, feature_digest=None, feature_dim=None)

Bases: ParameterEstimator

Fits inner on feature_fn(x) for raw items x – the estimator side of FeatureMapDensity.

Parameters:

• feature_name (str)

• inner (ParameterEstimator)

• name (str | None)

• feature_version (str | None)

• feature_digest (str | None)

• feature_dim (int | None)

accumulator_factory()

Return an accumulator factory that feature-maps raw inputs before inner accumulation.

Return type:

FeatureMapAccumulatorFactory

estimate(nobs, suff_stat)

Estimate the inner distribution and wrap it as a feature-map density.

Parameters:

• nobs (float | None)

• suff_stat (Any)

Return type:

FeatureMapDensity

feature_fn(name, *, version=None, digest=None, feature_dim=None)

Resolve a feature function and optionally verify its complete pinned identity.

Parameters:

• name (str)

• version (str | None)

• digest (str | None)

• feature_dim (int | None)

Return type:

Callable[[Any], ndarray]

register_feature_fn(name, fn, *, version, feature_dim)

Register one immutable, versioned raw-item feature transform.

Repeating the exact registration is idempotent. Reusing (name, version) for different code or output schema raises; callers must publish a new version instead.

Parameters:

• name (str)

• fn (Callable[[Any], ndarray])

• version (str)

• feature_dim (int)

Return type:

FeatureIdentity

class StreamingTransformer(module, device='cpu', *, lr=3.0e-3, optimizer_state=None)

Bases: SequenceEncodableProbabilityDistribution

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

Parameters:

• module (Any)

• device (str)

• lr (float)

• optimizer_state (dict[str, Any] | None)

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 next-token log probability for one (context, token) pair.

Parameters:

xy (Any)

Return type:

float

predict(x)

Return argmax next-token predictions for one or more contexts.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return the sampler for the conditional next-token model.

Parameters:

seed (int | None)

Return type:

StreamingTransformerSampler

seq_log_density(enc)

Return per-row next-token log probabilities for encoded context/token pairs.

Parameters:

enc (Any)

Return type:

ndarray

estimator(pseudo_count=None)

Return the streaming estimator that trains the live module in accumulator updates.

Parameters:

pseudo_count (float | None)

Return type:

StreamingTransformerEstimator

dist_to_encoder()

Return the encoder for context/token training pairs.

Return type:

StreamingTokenEncoder

to_dict()

Serialize the module bytes and device for registry-based round trips.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a StreamingTransformer from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

StreamingTransformer

NeuralLeaf

alias of NeuralGaussian

SoftmaxNeuralLeaf

alias of NeuralCategorical

StreamingTransformerLeaf

alias of StreamingTransformer

class TrainingSearchResult(recipe, loss, history=None, seed=None)

Bases: object

The outcome of a multi-fidelity training search: the best recipe, its full-budget loss, and the history.

The result OWNS what it reports (MXR-080-1890). It used to store the caller’s recipe mapping and history container by reference and validate nothing, so:

recipe = {"d_model": 64}
res = TrainingSearchResult(recipe=recipe, loss=1.0, history=trace)
recipe["d_model"] = 999          # the recorded winning recipe silently became a different one

rewrote a decision that had already been recorded, and loss accepted nan or the string "not a number" as a search outcome. __post_init__ now copies both containers with detach_receipt_container() and checks that the fields are the kind of thing the search actually produces.

seed records the search’s common-random-numbers seed, so the reported loss is attributable to a reproducible run rather than to one unlabelled draw.

Deliberately NOT checked, and why:

• The numeric contents of ``history``. detach_receipt_container copies containers but passes elements through by identity, so the X/Y blocks a search accumulates stay the same arrays rather than being duplicated – they are the optimizer’s own freshly built output, nobody else holds them, and copying every evaluation block would add a real cost to a result that is often only glanced at. A holder can still edit those arrays in place; this defends against the caller who kept a reference, which is the aliasing defect that was reproduced.

• That ``recipe`` decodes to ``history[“x”]``. Inside tune_training() both come from the same optimizer result, so the check would be tautological; for a directly constructed result there is no space to decode against.

Parameters:

• recipe (dict[str, Any])

• loss (float)

• history (Any)

• seed (int | None)

class TrainingSpace(d_model_choices=(64, 128, 256, 512), n_layer_range=(2, 12), log10_lr_range=(-4.0, -2.0), batch_choices=(16, 32, 64, 128))

Bases: object

The tunable axes of an LM training recipe and how a unit-cube point decodes into concrete knobs.

Parameters:

• d_model_choices (Sequence[int])

• n_layer_range (tuple[int, int])

• log10_lr_range (tuple[float, float])

• batch_choices (Sequence[int])

dims()

Return the dimensionality of the unit-cube recipe search space.

Return type:

int

bounds()

Return unit-cube bounds for the DOE optimizer.

Return type:

list[tuple[float, float]]

decode(point)

Decode a unit-cube point into concrete LM training hyperparameters.

Parameters:

point (ndarray)

Return type:

dict[str, Any]

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)

build_causal_lm(vocab, d_model=128, n_layer=3, n_head=4, block=64, embedding=None, gradient_checkpointing=False)

Build a causal decoder-only Transformer LM (token+pos embeddings, pre-norm blocks, weight-tied head).

embedding optionally injects a shared token nn.Embedding (vocab x d_model) to use in place of a fresh one – so several language models can tie the same word embedding and train it jointly (the weight-tied head follows it). Its shape must match (vocab, d_model).

gradient_checkpointing=True recomputes block activations during backward instead of storing them – identical gradients (pinned by test) for a large activation-memory cut on deep stacks or long blocks. The flag is a plain module attribute, so it can also be toggled on an existing model – including to a per-block list/tuple of bools (one per n_layer) rather than a single all-or-nothing bool, for F6’s cost-model-driven selective policy (mixle.models.memory_efficient_training.SelectiveRecomputePolicy).

Parameters:

• vocab (int)

• d_model (int)

• n_layer (int)

• n_head (int)

• block (int)

• embedding (Any)

• gradient_checkpointing (bool)

Return type:

Any

coarsen(model, budget, trust_region, input_law, n_mc=64, seed=0)

The iterated coarsening operator R: walk model.blocks pairwise, attempting a depth_merge() at each adjacent pair. A merge is ACCEPTED only if BOTH hold:

• TRUST REGION – its own LOCAL closed-form KL (receipt.kl_divergence) is at most trust_region;

• BUDGET – accepting it would not push the ACCUMULATED KL (summed over all accepted merges so far) past budget.

A rejected (or budget-exhausted) pair is left UNMERGED – both original blocks are kept, and the running law is propagated through them individually (via G1’s own per-layer laws, reusing _block_branch() plus the outer residual add) so later merge attempts still see the correct running law regardless of whether earlier pairs were merged. This makes the whole pass data-free: the running “receipt map” is built entirely from propagated LAWS, never real data.

Returns a CoarsenResult wrapping a new, real, forward-passable CoarsenedLM (so a caller can still measure REAL per-layer error against the original model by literally running both models on sampled token sequences – see mixle/tests/coarsening_test.py).

Parameters:

• model (Any)

• budget (float)

• trust_region (float)

• input_law (MultivariateGaussianDistribution)

• n_mc (int)

• seed (int)

Return type:

CoarsenResult

compress(model, method='auto', *, calibration_data=None, eval_data=None, sample_budget=None, input_law=None, budget=5.0, trust_region=5.0, n_mc=64, seed=0, kd_epochs=200, kd_lr=1e-2, hybrid_sample_fraction=0.01, hybrid_max_stages=1, hybrid_epochs=40, hybrid_lr=5e-4, target_n_layer=None)

The single compression entry point (roadmap J1): dispatches to one of the three method families, or lets a UCB1 CompressionReceipt-carrying picker choose (method="auto", default).

Parameters:

• model (Any) – a real mixle.models.transformer.CausalLM (or a coarsened one).

• method (str) – "auto" (default), "sampling_kd", "non_sampling", or "hybrid".

• calibration_data (Any) – integer token-context tensor/array (N, L) – required for "sampling_kd", "hybrid", and "auto" (the two methods that touch real data).

• eval_data (Any) – held-out integer token-context tensor/array used to MEASURE quality (teacher-agreement against model). Required and memory/value-independent from calibration_data for method="auto". Explicit methods may omit it, in which case quality is unmeasured rather than fabricated from training data.

• sample_budget (int | None) – an explicit cap (absolute count) on how many REAL calibration samples "hybrid"/"auto"’s hybrid arm may use; if None, hybrid_sample_fraction of len(calibration_data) is used instead.

• input_law (MultivariateGaussianDistribution | None) – the data-free input law for G1’s propagation; defaults to _default_input_law() (the model’s own token-embedding statistics).

• budget (float) – forwarded to mixle.models.coarsening.coarsen().

• trust_region (float) – forwarded to mixle.models.coarsening.coarsen().

• n_mc (int) – forwarded to mixle.models.coarsening.coarsen().

• kd_epochs (int) – forwarded to the fresh, from-scratch sampling_kd student’s full training loop.

• kd_lr (float) – forwarded to the fresh, from-scratch sampling_kd student’s full training loop.

• hybrid_sample_fraction (float) – fraction of calibration_data hybrid may use when sample_budget is not given (default 1%, matching the acceptance criterion).

• hybrid_max_stages (int) – how many worst-closure-error stages hybrid fine-tunes (default 1).

• hybrid_epochs (int) – forwarded to hybrid’s own micro-calibration fine-tune loop – deliberately separate from kd_epochs/kd_lr: hybrid fine-tunes ALREADY-good non_sampling weights on a HANDFUL of real samples (a light nudge), a completely different regime from training a fresh architecture from scratch on the full calibration set, and reusing the same hyperparameters for both badly overfits hybrid’s tiny sample budget in practice.

• hybrid_lr (float) – forwarded to hybrid’s own micro-calibration fine-tune loop – deliberately separate from kd_epochs/kd_lr: hybrid fine-tunes ALREADY-good non_sampling weights on a HANDFUL of real samples (a light nudge), a completely different regime from training a fresh architecture from scratch on the full calibration set, and reusing the same hyperparameters for both badly overfits hybrid’s tiny sample budget in practice.

• target_n_layer (int | None) – depth of the fresh sampling_kd student; defaults to coarsen’s own natural 2x depth cut (max(1, model.n_layer // 2)) so all three methods are compared at a MATCHED compression ratio.

• seed (int)

Returns:

CompressedModel

Return type:

CompressedModel

depth_merge(block_a, block_b, input_law, n_mc=64, seed=0)

Fold two adjacent Blocks into one via the second-order Taylor composition documented at module level.

Returns (merged_block, receipt) where merged_block is a real, forward-passable MergedBlock and receipt is a ScaleReceipt whose teacher_law/student_law are the EXACT-per-G1 sequential composition (block_a then block_b, propagated exactly as mixle.models.moment_propagation.propagate_moments() would) vs. the second-order MERGED composition, both Gaussian, so kl_divergence is the closed-form gaussian_kl() between them – the receipt for this individual (local) merge step, i.e. what a caller’s TRUST REGION check compares against.

Parameters:

• block_a (Any)

• block_b (Any)

• input_law (MultivariateGaussianDistribution)

• n_mc (int)

• seed (int)

Return type:

tuple[Any, ScaleReceipt]

experimental_width_merge_representation(target_width, input_law, temperature=0.1, n_iter=200)

Reduce the residual-stream width from d_model to target_width by pairing near-duplicate coordinates of the (Sigma-weighted) residual-stream covariance and merging/averaging them.

sigma_weighted_permutation (G2) was checked first per the roadmap note (see the module docstring): it solves min_P tr((W - P @ target_profile) Sigma (W - P @ target_profile)^T) for a SQUARE permutation P matching two SAME-shape objects (W against a fixed target_profile) – the classic one-to-one “git re-basin” alignment. Width reduction needs a genuinely MANY-TO-FEW map (d_model -> target_width, generally target_width < d_model so there is no permutation at all, square or otherwise), so it is not directly reusable here; _rectangular_sinkhorn() reuses the SAME log-domain Sinkhorn fixed-point idea for the rectangular marginals this problem actually has, rather than pulling in a separate heavy OT solver.

Data-free: the only input is input_law.covar (the propagated residual-stream covariance from G1), used to build a correlation-distance cost cost[i, j] = Sigma[i,i] + Sigma[a_j,a_j] - 2*Sigma[i, a_j] between every source coordinate i and target_width anchor coordinates a_j (the highest-variance coordinates, chosen as informative anchors) – cost[i, j] is exactly Var(x_i - x_{a_j}), so a near-zero cost means coordinate i is functionally redundant with anchor a_j and should be merged into it. The resulting Sinkhorn plan, column-normalized into convex combinations, is the merge operator; its pseudo-inverse is the reconstruction (“unmerge”) map.

This is explicitly a representation helper, not a model transformation: it does not rewire transformer weights, normalization, attention heads, or output heads and therefore does not return or claim a smaller executable model.

receipt.kl_divergence reports the true support result. It is infinite for a genuine rank reduction because the full-rank teacher assigns mass outside the rank-deficient round-trip support. For callers that explicitly assume isotropic downstream noise, regularized_kl_divergence and the receipt’s sensitivity values report that separate surrogate.

Parameters:

• target_width (int)

• input_law (MultivariateGaussianDistribution)

• temperature (float)

• n_iter (int)

Return type:

tuple[WidthMergeRepresentation, ScaleReceipt]

gaussian_kl(p, q)

Closed-form KL(p || q) for two multivariate Gaussians – the standard textbook formula:

KL(p||q) = 0.5 * ( tr(Sigma_q^-1 Sigma_p) + (mu_q - mu_p)^T Sigma_q^-1 (mu_q - mu_p)
                      - k + ln(det Sigma_q / det Sigma_p) )

computed ANALYTICALLY, not via Monte Carlo – both p (the “teacher” law) and q (the “student” law) are already MultivariateGaussianDistribution objects, which cache inv_covar and log_det from a (self-healing) Cholesky factorization at construction time, so this reuses those cached quantities directly rather than re-deriving them. Clipped at 0 to absorb float round-off on (near-)identical laws (KL is exactly 0 there, mathematically).

Parameters:

• p (MultivariateGaussianDistribution)

• q (MultivariateGaussianDistribution)

Return type:

float

structure_project(weight, sigma, mode='low_rank', rank=None, pattern='2:4')

Thin wrapper calling G2’s mixle.models.sigma_weighted_projection solvers directly – NOT a reimplementation, per the roadmap’s build-vs-borrow note. mode="low_rank" calls sigma_weighted_low_rank() (requires rank); mode="block_sparse" calls sigma_weighted_block_sparse() (uses pattern, either the literal "2:4" or an explicit boolean mask, exactly as G2 documents).

Parameters:

• weight (Any)

• sigma (Any)

• mode (str)

• rank (int | None)

• pattern (Any)

Return type:

tuple[ndarray, ProjectionReceipt]

evaluate_checkpoint(model, *, checkpoint_id='checkpoint', seed=0, n_examples=256, metadata=None, eval_data=None)

Run the full synthetic capability suite on model and return one structured EvalReport.

model is a real (or toy) mixle.models.transformer.CausalLM – anything exposing .vocab, .block, and __call__(x) -> (batch, vocab) next-token logits works. This is the single entry point: no separate setup per task, no manual orchestration – the “one command per checkpoint” F10 asks for.

Parameters:

• model (Any)

• checkpoint_id (str)

• seed (int)

• n_examples (int)

• metadata (dict[str, Any] | None)

• eval_data (Any)

Return type:

EvalReport

markov_transition_matrix(vocab)

The fixed order-1 Markov chain the perplexity task scores against.

Keyed only on vocab (via a module-level constant seed), not on the caller’s sample seed – the benchmark distribution is a fixed property of the eval suite (like a real held-out benchmark), while the per-call seed only controls which samples from it are drawn for a given run. This is what lets a training loop legitimately learn this chain (it is a fixed, nameable distribution) while individual eval runs still draw fresh, unseen sequences from it.

Parameters:

vocab (int)

Return type:

ndarray

track_regression(reports, *, threshold=0.05, reference='best')

Flag metrics that regressed by more than threshold (relative) across a sequence of checkpoints.

reports is an ordered sequence – successive training rungs, or successive steps of a J2 compression ladder. For each task, each report (from the second onward) is compared against either the best score seen so far in the sequence (reference="best", the default – catches slow drift, not just one-step drops) or the immediately-prior report (reference="prior" – catches only step-to-step drops). A task is flagged when it moved worse than the reference by more than threshold as a fraction of the reference’s magnitude, direction-aware (a drop for accuracy-like metrics, a rise for perplexity-like metrics).

Parameters:

• reports (Sequence[EvalReport])

• threshold (float)

• reference (str)

Return type:

RegressionReport

choose_compression_method(tensor, *, min_size_for_g4=_DEFAULT_MIN_SIZE_FOR_G4, g4_top_k_fraction=_DEFAULT_G4_TOP_K_FRACTION, g4_gof_threshold=DEFAULT_GOF_THRESHOLD, tail_family=None)

Goodness-of-fit-based per-tensor method picker (the “one level down” method-picker pattern I1 uses for its own picker – mixle.task.bandit is a reasonable fit for a LEARNED, reward- driven picker, but the choice here is cheaper and just as principled as a fixed rule: G4 already computes a real KS goodness-of-fit receipt as part of fitting, so reusing THAT receipt directly is simpler than bolting on a bandit that would need to learn what the receipt already tells us for free).

Rule: for tensors at or above min_size_for_g4 (below which G4’s fixed per-tensor fitting overhead is not worth it), attempt a G4 fit; if it does not dense-fall-back on its own goodness-of-fit receipt, use G4 (the more accurate, more expensive path). Otherwise, use int8 (the cheap path) – callers should still check compress_moment()’s returned method for a possible further downgrade to "dense" if int8 itself proves untrustworthy for this specific tensor (see _DEFAULT_INT8_ADVERSARIAL_RELATIVE_ERROR).

Parameters:

• tensor (Any)

• min_size_for_g4 (int)

• g4_top_k_fraction (float)

• g4_gof_threshold (float)

• tail_family (Any)

Return type:

str

compress_moment(tensor, method='auto', *, min_size_for_g4=_DEFAULT_MIN_SIZE_FOR_G4, g4_top_k_fraction=_DEFAULT_G4_TOP_K_FRACTION, g4_gof_threshold=DEFAULT_GOF_THRESHOLD, int8_block_size=_DEFAULT_INT8_BLOCK_SIZE, int8_adversarial_relative_error=_DEFAULT_INT8_ADVERSARIAL_RELATIVE_ERROR, int8_flushed_fraction_threshold=_DEFAULT_INT8_FLUSHED_FRACTION_THRESHOLD, tail_family=None)

Compress one Adam moment tensor (m or v) via G4, int8, or dense storage.

Parameters:

• tensor (Any) – A torch tensor or numpy array (one Adam moment buffer, any shape).

• method (str) – "auto" (use choose_compression_method()), "g4", "int8", or "dense" to force a specific path. A forced "g4"/"int8" still honestly downgrades to "dense" if the chosen method’s own receipt (G4’s KS statistic, or int8’s reconstruction error) rejects the fit – this function never returns a silently bad compressed representation.

• min_size_for_g4 (int)

• g4_top_k_fraction (float)

• g4_gof_threshold (float)

• int8_block_size (int)

• int8_adversarial_relative_error (float)

• int8_flushed_fraction_threshold (float)

• tail_family (Any)

Returns:

CompressedMomentEncoding

Return type:

CompressedMomentEncoding

decompress_moment(encoding)

Invert a complete versioned compressed-moment payload.

Parameters:

encoding (CompressedMomentEncoding)

Return type:

ndarray

dequantize_int8_blockwise(codes, scales, block_size=_DEFAULT_INT8_BLOCK_SIZE)

Invert quantize_int8_blockwise(). Returns a float32 array the same length as codes.

Parameters:

• codes (ndarray)

• scales (ndarray)

• block_size (int)

Return type:

ndarray

estimate_block_activation_bytes(batch, seq_len, d_model, dtype_bytes=4)

Estimate the memory footprint of ONE transformer block’s stored output activation (mixle.models.transformer.Block’s output, shape (batch, seq_len, d_model)) – the memory that activation checkpointing (recomputing instead of storing) frees.

Parameters:

• batch (int)

• seq_len (int)

• d_model (int)

• dtype_bytes (int)

Return type:

float

estimate_block_recompute_flops(batch, seq_len, d_model)

Estimate the FLOP cost of recomputing ONE transformer block’s forward pass, tied directly to mixle.models.transformer.Block’s actual layer shapes: qkv (d -> 3d), proj (d -> d), and the two MLP linears (d -> 4d -> d) give 3d^2 + d^2 + 4d^2 + 4d^2 = 12*d_model^2 linear-layer parameters per block; the standard “2 FLOPs per parameter per token” forward-pass heuristic turns that into a FLOP estimate, plus the attention score/value matmuls (QK^T and attn @ V, each ~2*batch*seq_len^2*d_model FLOPs) that scale with seq_len^2 rather than with parameter count.

Parameters:

• batch (int)

• seq_len (int)

• d_model (int)

Return type:

float

estimate_recompute_benefit(activation_bytes, memory_value_per_byte=_DEFAULT_MEMORY_VALUE_PER_BYTE)

The value of the memory freed by recomputing (rather than storing) one block’s activation.

Parameters:

• activation_bytes (float)

• memory_value_per_byte (float)

Return type:

float

estimate_recompute_cost(recompute_flops, flop_cost_per_unit=_DEFAULT_FLOP_COST_PER_UNIT)

The cost of the extra compute spent recomputing one block’s activation during backward.

Parameters:

• recompute_flops (float)

• flop_cost_per_unit (float)

Return type:

float

fp8_cast_with_guard(tensor, fp8_dtype='float8_e4m3fn', fallback_dtype=None, underflow_fraction_threshold=_DEFAULT_UNDERFLOW_FRACTION_THRESHOLD)

Attempt an fp8 cast of tensor, guarded against the two ways fp8 silently corrupts data.

Parameters:

• tensor (Any) – A torch tensor (any float dtype).

• fp8_dtype (str) – "float8_e4m3fn" (higher precision, smaller range – the default, matching common fp8-training practice for weights/activations) or "float8_e5m2" (wider range, coarser precision).

• fallback_dtype (Any) – Dtype to fall back to when a guard fires. Defaults to torch.bfloat16 (the codebase’s existing wide-range training dtype, per mixle.utils.parallel.torch_neural).

• underflow_fraction_threshold (float) – Maximum tolerable fraction of nonzero input values that are allowed to flush to exactly zero under the fp8 round-trip before the underflow guard fires.

Returns:

always carries the real, computed guard statistics, whether or not the fp8 cast was ultimately accepted – consistent with this codebase’s receipt-over-silent- assumption convention (see mixle.models.sorted_profile_quantizer’s goodness-of-fit receipt).

Return type:

Fp8CastResult

quantize_int8_blockwise(flat, block_size=_DEFAULT_INT8_BLOCK_SIZE)

Dynamic per-block symmetric int8 quantization (bitsandbytes-style 8-bit-Adam technique).

Each contiguous block of block_size elements gets its own scale (absmax / 127), so one extreme value only degrades the resolution of ITS OWN block rather than the whole tensor.

Returns:

(codes, scales) – codes is int8 and the same length as flat; scales is one float32 per block.

Return type:

tuple[np.ndarray, np.ndarray]

Parameters:

• flat (ndarray)

• block_size (int)

attention_law(law, qkv_weight, qkv_bias, proj_weight, proj_bias, n_head)

Propagate a Gaussian law through one mixle.models.transformer.CausalAttention layer.

Modeling assumption: the key/value population attended over (across sequence positions) is treated as a single STATIONARY Gaussian – the same joint law as the query’s – rather than tracking per-position laws. This is the “R2 MGF” population-level approximation the roadmap specifies; it does not model the causal mask explicitly (a real causal mask makes early positions attend to a smaller, non-stationary population). That mismatch is a known limitation, checked directly against Monte Carlo softmax attention in mixle/tests/moment_propagation_test.py (tightest for a roughly-stationary population, looser for strongly non-stationary / short sequences).

Derivation

For jointly Gaussian (K, V) (a key vector and its associated value vector from the SAME token) and a fixed query q, the requested MGF identity is

E[exp(q^T K / sqrt(d)) V] = exp(q^T mu_K/sqrt(d) + 0.5 q^T Sigma_KK q / d) * (mu_V + Sigma_VK q / sqrt(d))

which is the standard “E[Y e^{t^T X}] = M_X(t) (mu_Y + Sigma_YX t)” identity for jointly Gaussian (X, Y) with t = q / sqrt(d), X = K, Y = V. The attention DENOMINATOR (the softmax normalizer) is exactly the same MGF evaluated with V replaced by the constant 1:

E[exp(q^T K / sqrt(d))] = exp(q^T mu_K/sqrt(d) + 0.5 q^T Sigma_KK q / d).

Both share the identical exponential prefactor, so it CANCELS in the ratio:

softmax-attention-output(q) ~= E[exp(q^T K/sqrt(d)) V] / E[exp(q^T K/sqrt(d))] = mu_V + (Sigma_VK / sqrt(d)) q.

This is a remarkably clean result: the MGF-approximated attention output, as a function of the query, is EXACTLY AFFINE in q (no exponential term survives). Since q itself has a propagated Gaussian law (the marginal of the joint (Q, K, V) law after the qkv projection), pushing that law through this affine map is exact (reusing linear_law()’s formula) – there is no additional approximation beyond the population-stationarity assumption above and (for the covariance) the affine relation being evaluated with the Sigma_VK estimated from the SAME joint law used for the mean.

Per head, per token position, y_h(q) = mu_{V,h} + (Sigma_{VK,h} / sqrt(d_head)) q_h. Stacking heads into a block-diagonal map A (each head only reads its own query slice) and reusing the FULL (cross-head) query covariance from the qkv projection gives the cross-head covariance of the output “for free” – cross-head correlation in Q (already present in Sigma_QQ’s off-diagonal blocks) propagates through A Sigma_QQ A^T even though A itself only mixes within a head.

Parameters:

• law (MultivariateGaussianDistribution)

• qkv_weight (ndarray)

• qkv_bias (ndarray | None)

• proj_weight (ndarray)

• proj_bias (ndarray | None)

• n_head (int)

Return type:

tuple[MultivariateGaussianDistribution, ndarray]

gelu_law(law, *, return_receipt=False)

Propagate a Gaussian law through elementwise GELU using the closed-form per-dimension moments in _gelu_scalar_moments().

The per-dimension MEAN and VARIANCE are exact closed-form expressions (no Monte Carlo, no approximation of the GELU functional form). The OFF-diagonal output covariance – Cov(GELU(x_i), GELU(x_j)) for i != j – has no simple closed form (it needs the joint bivariate distribution of (x_i, x_j) through a nonlinearity) and is instead approximated by the standard delta-method / Price’s-theorem-style linearization Cov(y_i, y_j) ~= J_ii * J_jj * Cov(x_i, x_j) with J_ii = d E[GELU(x_i)]/d mu_i the EXACT mean-derivative from Stein’s lemma. The diagonal is then overwritten with the exact closed-form variance. Because that splice need not be PSD, correlations are projected to the PSD cone under a fixed-unit-diagonal constraint and then rescaled by the exact marginal standard deviations. This preserves the scalar variances while correcting only the approximate correlation structure.

Set return_receipt=True to receive a GeluCovarianceReceipt measuring the correction, pre/post minimum correlation eigenvalues, iteration count, and marginal-variance preservation error.

Parameters:

• law (MultivariateGaussianDistribution)

• return_receipt (bool)

Return type:

tuple[MultivariateGaussianDistribution, ndarray] | tuple[MultivariateGaussianDistribution, ndarray, GeluCovarianceReceipt]

iter_moments(model, input_law, n_mc=128, seq_len=16, seed=0, probe_batch_size=32)

Yield layer-local Gaussian laws without retaining earlier receipts.

input_law models the distribution of a token’s residual-stream vector ENTERING the block stack (i.e. after token + position embedding) – an N(mu, Sigma) over R^{d_model}.

The iterator walks model.blocks one at a time, then the final model.ln and model.head. Closure probes execute in the target module’s device/dtype and evaluation mode, restore every submodule’s prior mode, and reduce batches of at most probe_batch_size online.

Parameters:

• model (Any)

• input_law (MultivariateGaussianDistribution)

• n_mc (int)

• seq_len (int)

• seed (int)

• probe_batch_size (int)

Return type:

Iterator[LayerMoments]

layernorm_law(law, weight, bias, eps=1e-5)

Propagate a Gaussian law through LayerNorm: y = weight * (x - m(x)) / sqrt(v(x) + eps) + bias, where m(x) = mean_d(x) and v(x) = mean_d((x - m(x))^2) are the PER-SAMPLE (per-token) statistics LayerNorm computes over the feature axis – this is the nonlinear, data-dependent step in the block.

Derivation (“re-anchoring”)

LayerNorm has no closed-form pushforward of a full law in general (m(x) and v(x) are themselves random, nonlinear functions of x). We use a first-order Taylor expansion around the CURRENT mean mu, i.e. anchor the (unknown, per-sample) normalization statistics at their EXPECTED values under the current law, not merely their values evaluated at the mean vector. m(x) = mean_d(x) is itself LINEAR in x, so E[m(x)] = mean_d(mu) exactly – no correction needed there. But v(x) = mean_d((x - m(x))^2) is QUADRATIC in x, so evaluating it at mu alone (v(mu) = mean((mu - m)^2)) drops a systematic bias term: writing P = I - (1/d) 11^T for the feature-centering projector (symmetric, idempotent), v(x) = (1/d) x^T P x, and the standard quadratic-form expectation identity gives

E[v(x)] = v(mu) + (1/d) trace(P @ Sigma) = v(mu) + (1/d) (trace(Sigma) - (1/d) sum(Sigma)).

The second term is the “spread of x around its own per-token mean” contribution that v(mu) alone misses entirely; it is NOT a higher-order correction that can be dropped once Sigma is non-negligible relative to d – for small d_model (e.g. an 8-wide toy model) it routinely dominates v(mu), which without this correction makes the anchored std = sqrt(v(mu) + eps) far too small and blows the propagated mean/covariance up by an order of magnitude relative to the true LayerNorm output law. We therefore re-anchor at the BIAS-CORRECTED v = E[v(x)] above (still a first-order/delta-method treatment of the covariance push-forward – only the anchor point for v is corrected, not the linearization itself). This is the LLM analogue of BatchNorm-based data-free-quantization (DFQ) calibration, which likewise re-derives cheap closed-form layer statistics (there, running mean/var; here, the expected LayerNorm mean/var under the CURRENT propagated law) without touching real data.

The mean is propagated by evaluating the true (nonlinear) LayerNorm map at mu but with the bias-corrected v:

mu_out = weight * (mu - m) / sqrt(v + eps) + bias.

The covariance is propagated via the JACOBIAN of LayerNorm evaluated at mu with the same bias-corrected v (a standard, textbook LayerNorm-backward-style derivative applied at the corrected anchor):

d(norm_i)/d(x_j) |_{x=mu} = (1/sqrt(v+eps)) * (delta_ij - 1/d - (mu_i - m)(mu_j - m) / (d*(v+eps))) J_ij = weight_i * d(norm_i)/d(x_j) Sigma_out ~= J Sigma J^T.

Known failure modes (feeds the closure-error receipt)

• Small ``d_model``: m(x) and v(x) are averages over only d_model features, so their sample-to-sample fluctuation around their (now bias-corrected) expectation – which this delta-method approximation still ignores, since the Jacobian itself is frozen at the anchor – is large relative to their magnitude when d_model is small. The approximation is progressively worse as d_model shrinks, even after the mean-bias correction above.

• Heavy-tailed pre-norm activations: the first-order Taylor expansion is only locally valid; if the input law has high kurtosis / is far from Gaussian in practice (despite being MODELED as Gaussian here), the true per-sample (m, v) can swing far from their Gaussian-law expectations, and the linearization degrades.

• A deep stack of blocks compounds both effects: even a small per-layer LayerNorm error can accumulate across n_layer re-anchoring steps.

Parameters:

• law (MultivariateGaussianDistribution)

• weight (ndarray)

• bias (ndarray)

• eps (float)

Return type:

tuple[MultivariateGaussianDistribution, ndarray]

linear_law(law, weight, bias=None)

Exact Gaussian push-forward of y = Wx + b.

x ~ N(mu, Sigma) => y ~ N(W mu + b, W Sigma W^T) exactly – no approximation. Returns the new law together with the Jacobian dy/dx = W (used to chain residual cross-covariances).

Note on out.covar when weight maps to MORE output dimensions than law.covar’s rank (e.g. an expanding projection like the qkv weight in attention_law()): W Sigma W^T is then mathematically rank-deficient (PSD, not PD), and the returned law – a MultivariateGaussianDistribution, which self-heals a non-PD covariance with a small trace-scaled diagonal jitter so it stays Cholesky-factorizable for log_det/inv_covar/sampling, see _robust_cho_factor in mixle.stats.multivariate.multivariate_gaussian – may therefore differ from the exact W Sigma W^T by that jitter (diagonal-only, ~1e-10 * trace/n). This is NOT an error in the formula above – verified bit-exact whenever the output is genuinely PD (e.g. square invertible weight); see mixle/tests/moment_propagation_test.py::LinearExactTest for both regimes – it is the unavoidable cost of representing an exactly-singular covariance with a factorization that requires one.

Parameters:

• law (MultivariateGaussianDistribution)

• weight (ndarray)

• bias (ndarray | None)

Return type:

tuple[MultivariateGaussianDistribution, ndarray]

propagate_moments(model, input_law, n_mc=128, seq_len=16, seed=0, probe_batch_size=32)

Materialize iter_moments() into a list of per-layer full-covariance receipts.

Unlike the iterator, this convenience API intentionally retains O(depth × d_model²) returned law state. It remains useful when callers need random access to every layer; callers requiring depth-independent retained memory should consume iter_moments().

Parameters:

• model (Any)

• input_law (MultivariateGaussianDistribution)

• n_mc (int)

• seq_len (int)

• seed (int)

• probe_batch_size (int)

Return type:

list[LayerMoments]

build_auto_optimizer(module, *, lr=1.0e-3, sign_stable=True, **kwargs)

Plan and build the default non-Adam-first optimizer for module.

Parameters:

• module (Any)

• lr (float)

• sign_stable (bool)

• kwargs (Any)

Return type:

Any

build_routed_optimizer(module, plan, *, lr=1.0e-3, betas=(0.9, 0.999), eps=1.0e-8, weight_decay=0.0, precondition_frequency=10, muon_backend='newton_schulz', muon_steps=7, fisher_provider=None, projection_provider=None)

Build a Torch optimizer from stable or experimental structurally compatible routes.

Parameters:

• module (Any)

• plan (Any)

• lr (float)

• betas (tuple[float, float])

• eps (float)

• weight_decay (float)

• precondition_frequency (int)

• muon_backend (str)

• muon_steps (int)

• fisher_provider (Callable[[str, Any], Any] | None)

• projection_provider (Callable[[str, Any], None] | None)

Return type:

Any

plan_neural_optimizer(module, *, matrix_min_elements=4_096, sign_stable=True)

Plan update families from parameter role and shape without importing Torch at module import time.

Parameters:

• module (Any)

• matrix_min_elements (int)

• sign_stable (bool)

Return type:

NeuralOptimizerPlan

resolve_neural_optimizer(module, optimizer=None, *, lr=1.0e-3, sign_stable=True)

Resolve automatic, named, planned, or callable neural optimizers with an audit receipt.

Parameters:

• module (Any)

• optimizer (Any)

• lr (float)

• sign_stable (bool)

Return type:

tuple[Any, dict[str, Any]]

shard_safe_neural_optimizer_plan(plan)

Localize routes whose geometry is not invariant under parameter sharding.

Muon/Kronecker/natural-gradient directions require a global matrix. Applying them independently to rectangular shards is a different algorithm. Until a backend declares and executes the required gather/factor reductions, use momentum for those blocks and record the reason in the plan.

Parameters:

plan (NeuralOptimizerPlan)

Return type:

NeuralOptimizerPlan

consistency_loss(student_output, teacher_output, mode='mse')

The self-distillation consistency term between a student prediction and a teacher/target prediction on the SAME input – mode="mse" (default, plain squared-error between logits, the mean-teacher convention) or mode="kl" (KL(teacher_softmax || student_log_softmax), the classic soft-target distillation loss).

Parameters:

• student_output (Any)

• teacher_output (Any)

• mode (str)

Return type:

Any

stochastic_depth_forward(model, x, drop_prob, generator=None, *, return_keep_mask=False)

Run model on x twice: once at full depth, once with each block independently dropped with probability drop_prob (at least one block is always kept, so the partial pass never degenerates to the bare embedding/head). Returns (full_output, partial_output) – the pair train_with_self_distillation() feeds to consistency_loss().

At drop_prob == 0 both passes keep every block, so the two outputs are IDENTICAL (no dropout elsewhere in Block) – the degenerate-case sanity check pinned in mixle/tests/self_distillation_test.py.

Parameters:

• model (Any)

• x (Any)

• drop_prob (float)

• generator (Any)

• return_keep_mask (bool)

Return type:

tuple

train_with_self_distillation(model, data, steps, *, ema_decay=0.999, drop_prob=0.1, consistency_weight=1.0, ema_weight=None, stochastic_depth_weight=None, consistency_mode='mse', lr=3e-3, device='cpu', optimizer=None, seed=0, log=None)

Train model (a CausalLM, trained in place and also returned) for steps next-token cross-entropy steps, with EMA-teacher consistency and stochastic-depth consistency added as extra loss terms on top of the SAME per-step batch – both self-distillation pressures happen DURING training, not as a separate post-hoc pass.

data yields (context, next_token) micro-batches shaped exactly like mixle.data.stream_token_source.stream_token_source() (context: (batch, block) float ids, next_token: (batch,) int ids). data may be:

• a zero-arg CALLABLE returning a fresh iterator each time (e.g. lambda: stream_token_source(ids, block=64, batch_size=32)) – restarted automatically whenever it runs dry before steps is reached, so training can outlast one epoch; or

• a plain iterable/iterator (e.g. a list of batches, or a single generator object) – consumed once, sized to yield at least steps batches (a bare generator can’t be rewound).

Per step: loss = cross_entropy(full_depth_logits, target) + stochastic_depth_weight * consistency(partial_depth_logits, full_depth_logits.detach()) + ema_weight * consistency(full_depth_logits, ema_teacher(context)), then one optimizer step, then one EMA-teacher update. ema_weight/stochastic_depth_weight each default to consistency_weight when unset.

Parameters:

• model (Any)

• data (Any)

• steps (int)

• ema_decay (float)

• drop_prob (float)

• consistency_weight (float)

• ema_weight (float | None)

• stochastic_depth_weight (float | None)

• consistency_mode (str)

• lr (float)

• device (str)

• optimizer (Any)

• seed (int)

• log (Any)

Return type:

Any

ewc(anchor, fisher, lam=1.0, floor=0.0)

Bundle (anchor, fisher, lambda) for .fit(..., ewc=...) (the EWC anti-forgetting penalty).

floor adds a constant to every Fisher entry, so the penalty weights are F_i + floor and the penalty is lambda * sum_i (F_i + floor) (theta_i - theta*_i)^2. It defaults to 0.0, which is EWC exactly as defined – and which, on its own, usually does not retain the old task, for a reason that is structural rather than a matter of tuning lam.

The diagonal empirical Fisher is estimated from per-example gradients at the end of the previous stage. A stage that fit well leaves those gradients near zero, so almost every entry of F is near zero and the quadratic it defines is badly rank-deficient. Minimizing that penalty then does not hold the module near theta*; it moves the module through the Fisher’s null space, where the penalty is cheap and the old task’s function is not preserved. Measured on a 4-feature 3-class pair of unrelated tasks (see neural_ppl_test.test_cpt_ewc_retains_the_old_task), raising lam from 10 to 8000 drives the Fisher-weighted drift down by four orders of magnitude, 3.6e-1 to 3.7e-5, while the raw distance from the anchor nearly doubles, 81 to 147, and retained accuracy on the old task stays flat around 0.21-0.24 against 0.19 for plain continuation. The penalty is satisfied; the capability it exists to protect is not. Past lam of about 1e4 the quadratic is stiffer than the routed optimizer’s step size and the objective diverges instead.

A small positive floor closes that null space – it is EWC plus L2-to-anchor, a Tikhonov regularization of the importance estimate rather than of the parameters. On the same task pair floor=1e-3, lam=1e2 retains 0.59-0.74 on the old task against 0.19-0.34 for plain continuation, at 0.62-0.66 on the new one: a real stability/plasticity trade rather than a satisfied constraint. Prefer a floor over a large lam; the floor buys retention where lam alone only buys stiffness.

Parameters:

• anchor (ParameterBundle)

• fisher (ParameterBundle)

• lam (float)

• floor (float)

Return type:

tuple

extrapolate_learning_curve(steps, losses, *, at)

Predict the loss at budget/step at from a partial run’s (steps, losses) via a power-law fit.

Fits loss(t) = a + b * t^(-c) and evaluates it at at so a partial run can estimate the full-budget loss for early stopping. Invalid data, an underdetermined curve, or a failed/non-finite fit raises explicitly.

Parameters:

• steps (Sequence[float])

• losses (Sequence[float])

• at (float)

Return type:

float

fisher_diagonal(leaf, x, y, *, samples=512, device='cpu', seed=0)

Diagonal empirical Fisher computed on an isolated evaluation-mode copy.

Parameters:

• leaf (Any)

• x (Any)

• y (Any)

• samples (int)

• device (str)

• seed (int)

Return type:

ParameterBundle

lm_train_fn(token_ids, val_ids, *, vocab, block=64, max_epochs=3, device='cpu', seed=0, fidelities=None)

Return a training callback (recipe, budget) -> held-out nats/token for LM.

budget in (0, 1] is a fraction of the FULL run’s training tokens (max_epochs * len(token_ids)), not a fraction of the epoch count. budget=1.0 is unchanged – max_epochs epochs over the whole corpus – while a smaller budget shortens the corpus prefix as well as the epoch count, so distinct budgets buy distinct work. See LMFidelity for the measurement that motivated this (MXR-080-1890).

seed is the common-random-numbers seed shared by every run, which is what makes two recipes comparable; see LMTrainFn.__call__().

fidelities is optional and off by default. Pass the exact budget ladder the search will use and it is checked, before any training happens, for budgets that still resolve to identical work – the failure this finding is about. Token resolution runs out on a short corpus, so this can legitimately happen; it is reported rather than assumed away.

Raises:

ValueError – If fidelities is supplied and two of its budgets execute the same work.

Parameters:

• token_ids (Sequence[int])

• val_ids (Sequence[int])

• vocab (int)

• block (int)

• max_epochs (int)

• device (str)

• seed (int)

• fidelities (Sequence[float] | None)

Return type:

LMTrainFn

snapshot(leaf_or_module)

Detached clones of the module’s parameters – the anchor theta* for an EWC penalty.

Parameters:

leaf_or_module (Any)

Return type:

ParameterBundle

tune_training(train, space=None, *, fidelities=(0.25, 1.0), costs=None, max_cost=20.0, n_init=None, seed=0)

Run multi-fidelity BO over a training recipe.

train(recipe, budget) returns held-out loss, where lower is better. fidelities are the training-budget fractions the search may run at. Returns the recipe with the best full-budget loss and the full BO history.

seed is recorded on the returned result (MXR-080-1890) so the reported loss is attributable to a reproducible search rather than to one unlabelled draw. Whether the objective is itself reproducible is the callback’s contract, not this function’s: lm_train_fn() supplies one that pins its own randomness, but an arbitrary caller-supplied train is opaque here and is deliberately not second-guessed.

Parameters:

• train (Callable[[dict[str, Any], float], float])

• space (TrainingSpace | None)

• fidelities (tuple[float, ...])

• costs (tuple[float, ...] | None)

• max_cost (float)

• n_init (int | None)

• seed (int)

Return type:

TrainingSearchResult

apply_mup_init(model, *, base_width, base_std=0.02)

Re-initialize model in place per the muP init rules, relative to base_width.

base_std is the hidden-role init std tuned/measured at base_width (the mixle default, 0.02, matches common transformer practice and is a reasonable base-width value on its own). model.d_model is read as the target width, so width_mult = model.d_model / base_width. LayerNorm weight/bias keep their identity init (1 / 0, unaffected by width, matching the "input" role’s no-rescale treatment); every other bias is zero-initialized (muP does not rescale bias init); every other weight matrix is drawn Normal(0, base_std * init_std_multiplier(role, width_mult)). Also turns on muP attention-logit scaling (enable_mup_attention()) on every block – the 1/head_dim QK scaling is as much a part of “the model is parametrized under muP” as the init/lr rules above, and previously being left at the standard 1/sqrt(head_dim) scale was an unintentional gap between what this module documented and what it actually configured.

Parameters:

• base_width (int)

• base_std (float)

Return type:

None

classify_causal_lm_params(model)

Map every named parameter of a mixle.models.transformer.CausalLM to its muP role.

• tok.weight / pos.weight (embeddings) -> "input" – fan-in is the fixed vocab / block length, not d_model.

• LayerNorm affine params (blocks.*.ln1/ln2, top-level ln) -> "input" – no fan-in at all, muP leaves them at their standard Theta(1) scale/shift regardless of width.

• head.weight -> not a separate entry: it is the same nn.Parameter as tok.weight (weight tying), so model.named_parameters() already reports it once, under "tok.weight". See the module docstring for how the muP output role is instead applied at readout time.

• everything else (the attention qkv/proj and MLP Linear weights/biases inside each block) -> "hidden" – fan-in scales linearly with d_model.

Return type:

dict[str, Literal[‘input’, ‘hidden’, ‘output’]]

init_std_multiplier(role, width_mult)

Return the muP multiplier on init std for a parameter of role at the given width_mult.

width_mult = target_width / base_width. Multiply the BASE width’s tuned/measured init std by this factor to get the init std to use at the target width:

• "input" -> 1 (fan-in fixed, e.g. vocab size – no rescale)

• "hidden" -> width_mult ** -0.5 (variance 1/fan_in, standard, unchanged form)

• "output" -> width_mult ** -1 (variance 1/fan_in**2 – an extra 1/width_mult)

Parameters:

• role (Literal['input', 'hidden', 'output'])

• width_mult (float)

Return type:

float

lr_multiplier(role, width_mult)

Return the muP multiplier on learning rate for a parameter of role at the given width_mult.

width_mult = target_width / base_width. Multiply the BASE width’s tuned lr by this factor to get the transferred lr to use at the target width:

• "input" -> 1 (constant lr – the muP “don’t touch it” role)

• "hidden" -> width_mult ** -1 (the headline muP rule: lr shrinks as the model widens)

• "output" -> width_mult ** -1 (same shrink as hidden, for Adam)

Parameters:

• role (Literal['input', 'hidden', 'output'])

• width_mult (float)

Return type:

float

mup_param_groups(model, *, base_width, lr)

Build torch optimizer param groups implementing muP’s per-role lr scaling for model.

lr is the BASE (hidden-role) learning rate – the one hyperparameter tuned once, cheaply, at base_width. model.d_model is read as the target width. Returns a list of {"params": [...], "lr": ..., "mup_role": ...} dicts suitable for torch.optim.Adam(groups); the "input" group’s lr is unscaled, the "hidden"/"output" groups get lr * lr_multiplier(role, width_mult) – i.e. passing the same tuned lr at any target width reproduces exactly what transfer_lr() predicts for that role. transfer_lr is the formula; this is the mechanism that applies it to a live model + optimizer.

Parameters:

• base_width (int)

• lr (float)

Return type:

list[dict]

output_forward_multiplier(width_mult)

Return the muP readout multiplier (the “c” of abc-parametrization) applied to output-role logits.

Multiply the raw head output by this factor so the readout’s output scale stays width-independent at init, even though (due to weight tying, see the module docstring) head.weight itself is parametrized under the "input" rule rather than a separate "output" init/lr rule.

Parameters:

width_mult (float)

Return type:

float

transfer_init_std(base_std, base_width, target_width, *, role='hidden')

Rescale base_std (tuned/measured at base_width) to the muP init std at target_width.

Parameters:

• base_std (float)

• base_width (int)

• target_width (int)

• role (Literal['input', 'hidden', 'output'])

Return type:

float

transfer_lr(base_lr, base_width, target_width, *, role='hidden')

Rescale base_lr (tuned at base_width) to the muP-predicted optimum at target_width.

This is the deliverable that “collapses the ladder’s tuning bill”: tune role="hidden" lr once, cheaply, at a small base_width, then call this to predict the optimal lr at any larger target_width with (ideally) no further search. Defaults to role="hidden" – the dominant parameter group (attention + MLP) and the one muP’s headline lr-transfer guarantee is about.

Parameters:

• base_lr (float)

• base_width (int)

• target_width (int)

• role (Literal['input', 'hidden', 'output'])

Return type:

float

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

Train one persistent-optimizer step per streamed micro-batch without buffering the corpus.

Parameters:

• module (Any)

• token_source (Any)

• lr (float)

• device (str)

• report_every (int)

• log (Any)

Return type:

tuple

class GrammarLearningResult(model, history, validation_history=None)

Bases: FitResult[HeterogeneousPCFGDistribution]

Fitted PCFG plus training and optional validation log-likelihood history.

Parameters:

• model (ModelT)

• history (list[float])

• validation_history (list[float] | None)

class HamiltonianNet(dim, hidden=(64, 64), *, separable=True)

Bases: object

Learns a scalar H(q, p) and exposes the symplectic-gradient dynamics it implies.

dim is the dimension of q (and of p, always equal); hidden sizes the network computing H. By default the network is masked into independent V(q) and T(p) branches whose scalar outputs are summed, making the supported leapfrog premise structural rather than advisory. Set separable=False only when learning a general Hamiltonian that will be integrated by another method. module (the underlying torch.nn.Module) is exposed directly for training with an ordinary optimizer loop against derivative-matching data.

Parameters:

• dim (int)

• hidden (Sequence[int])

• separable (bool)

hamiltonian(q, p)

H(q, p), shape (...,) – squeezes the module’s scalar output dimension.

Parameters:

• q (Any)

• p (Any)

Return type:

Any

time_derivative(q, p)

(dq/dt, dp/dt) = (dH/dp, -dH/dq) via autograd – the symplectic gradient of hamiltonian.

q/p must be leaf tensors (requires_grad_(True)) or already part of an active graph; the returned derivatives carry gradients back through self.module’s parameters, so this composes into a training loop that fits derivative-matching data end to end.

Parameters:

• q (Any)

• p (Any)

Return type:

tuple[Any, Any]

class HardEMResult(model, history, validation_history=None)

Bases: FitResult[StochasticBlockGraphModel]

Result from hard-EM fitting of a stochastic block model.

Parameters:

• model (ModelT)

• history (list[float])

• validation_history (list[float] | None)

class KnowledgeGraphFitResult(model, history, validation_history=None)

Bases: FitResult[TransEKnowledgeGraphModel]

Result from TransE margin fitting.

Parameters:

• model (ModelT)

• history (list[float])

• validation_history (list[float] | None)

class PartiallyObservableMarkovDecisionProcessFilterResult(beliefs, log_likelihood, predictive_observation_probs, possible=True, impossible_at=None)

Bases: object

Belief trajectories, log likelihood, and predictive observation terms.

Parameters:

• beliefs (ndarray)

• log_likelihood (float)

• predictive_observation_probs (ndarray)

• possible (bool)

• impossible_at (int | None)

class PartiallyObservableMarkovDecisionProcessFitResult(model, history, validation_history=None)

Bases: FitResult[PartiallyObservableMarkovDecisionProcessModel]

Baum-Welch style fit result for known-action PartiallyObservableMarkovDecisionProcess sequences.

Parameters:

• model (ModelT)

• history (list[float])

• validation_history (list[float] | None)

class PartiallyObservableMarkovDecisionProcessModel(transition, observation, initial_belief=None, rewards=None, name=None)

Bases: object

Finite-state PartiallyObservableMarkovDecisionProcess with action-conditioned transitions and observations.

transition[a, i, j] is P(S_t=j | S_{t-1}=i, A_t=a). observation[a, j, o] is P(O_t=o | S_t=j, A_t=a).

Parameters:

• transition (Any)

• observation (Any)

• initial_belief (Any | None)

• rewards (Any | None)

• name (str | None)

belief_update(belief, action, observation)

Update a belief after taking action and seeing observation.

Parameters:

• belief (Any)

• action (int)

• observation (int)

Return type:

tuple[ndarray, float]

filter(actions, observations, initial_belief=None)

Run the forward filter and return posterior beliefs and log likelihood.

Parameters:

• actions (Sequence[int])

• observations (Sequence[int])

• initial_belief (Any | None)

Return type:

PartiallyObservableMarkovDecisionProcessFilterResult

sequence_log_likelihood(actions, observations, initial_belief=None)

Return log P(observations | actions, model).

Parameters:

• actions (Sequence[int])

• observations (Sequence[int])

• initial_belief (Any | None)

Return type:

float

forward_backward(actions, observations, initial_belief=None)

Return state marginals, transition marginals, and sequence log likelihood.

Parameters:

• actions (Sequence[int])

• observations (Sequence[int])

• initial_belief (Any | None)

Return type:

tuple[ndarray, ndarray, float]

predict_observation(belief, action)

Return P(O_t | belief, action) before observing O_t.

Parameters:

• belief (Any)

• action (int)

Return type:

ndarray

expected_reward(belief, action)

Return E[R | belief, action] when rewards were supplied.

Parameters:

• belief (Any)

• action (int)

Return type:

float

sample(actions, seed=None, initial_belief=None)

Sample latent states and observations for a fixed action sequence.

Parameters:

• actions (Sequence[int])

• seed (int | None)

• initial_belief (Any | None)

Return type:

tuple[ndarray, ndarray]

class PartiallyDirectedGraph(directed_edges, undirected_edges, variable_names)

Bases: object

Partially directed graph after collider orientation.

Parameters:

• directed_edges (set[tuple[int, int]])

• undirected_edges (set[tuple[int, int]])

• variable_names (list[Any])

class PCFGParseNode(label, span, log_prob, rule_index, rule_type, children=(), value=None)

Bases: object

Node in a Viterbi parse tree.

Parameters:

• label (Any)

• span (tuple[int, int])

• log_prob (float)

• rule_index (int)

• rule_type (str)

• children (tuple[PCFGParseNode, ...])

• value (Any)

leaves()

Return terminal observations under this node.

Return type:

list[Any]

class PINNRegression(module, residual_fn, domain, *, problem, noise=1.0, residual_weight=1.0, n_collocation=64, m_steps=40, lr=0.01, seed=0, name=None, device=None, constraint_receipt=None)

Bases: NeuralGaussian

NeuralGaussian plus a PDE/ODE-residual penalty evaluated on sampled collocation points.

domain is a (low, high) pair of per-dimension box bounds for collocation sampling; residual_fn computes the physics residual (see module docstring); residual_weight scales the penalty relative to the data-fit NLL; n_collocation is how many collocation points are drawn fresh every M-step.

Parameters:

• module (Any)

• residual_fn (Any)

• domain (tuple[Any, Any])

• problem (PINNProblemSpec)

• noise (float)

• residual_weight (float)

• n_collocation (int)

• m_steps (int)

• lr (float)

• seed (int)

• name (str | None)

• device (Any)

• constraint_receipt (PINNConstraintReceipt | None)

estimator(pseudo_count=None)

Return the estimator that combines weighted data fit with residual collocation penalties.

Parameters:

pseudo_count (float | None)

Return type:

PINNRegressionEstimator

dist_to_encoder()

Return the neural-Gaussian encoder for (x, y) observation pairs.

Return type:

NeuralGaussianEncoder

to_dict()

Serialize the module, residual function reference, domain, and PINN hyperparameters.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a PINNRegression from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

PINNRegression

class PINNRegressionEstimator(module, residual_fn, domain, *, problem, noise=1.0, residual_weight=1.0, n_collocation=64, m_steps=40, lr=0.01, seed=0, name=None, device=None)

Bases: NeuralGaussianEstimator

EM estimator for PINNRegression: the M-step adds a residual penalty on fresh collocation points to the same weighted-NLL gradient descent NeuralGaussianEstimator runs.

Collocation sampling is deterministic given seed (a private numpy.random.RandomState, advanced once per M-step) – refitting with the same seed draws the same collocation batches.

Parameters:

• module (Any)

• residual_fn (Any)

• domain (tuple[np.ndarray, np.ndarray])

• problem (PINNProblemSpec)

• noise (float)

• residual_weight (float)

• n_collocation (int)

• m_steps (int)

• lr (float)

• seed (int)

• name (str | None)

• device (Any)

accumulator_factory()

Return the neural-Gaussian accumulator factory for weighted observation pairs.

Return type:

NeuralGaussianAccumulatorFactory

estimate(nobs, suff_stat)

Run the data-plus-residual M-step and return the updated PINN leaf.

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

PINNRegression

class PINNProblemSpec(equation, coordinates, fields, residual_components, constraints, identifiability_conditions, constraint_enforcement, completeness_basis)

Bases: object

Explicit mathematical and enforcement schema for a PINN fit.

Parameters:

• equation (str)

• coordinates (tuple[str, ...])

• fields (tuple[str, ...])

• residual_components (tuple[str, ...])

• constraints (tuple[str, ...])

• identifiability_conditions (tuple[str, ...])

• constraint_enforcement (str)

• completeness_basis (str)

to_dict()

Return the JSON-compatible problem declaration.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Validate and restore a problem declaration.

Parameters:

payload (Any)

Return type:

PINNProblemSpec

class PINNConstraintReceipt(status, equation, constraint_enforcement, constraints, identifiability_conditions, completeness_basis, labeled_constraints, residual_loss_initial, residual_loss_final)

Bases: object

Durable declaration and measured outcome of a constraint-aware PINN fit.

Parameters:

• status (str)

• equation (str)

• constraint_enforcement (str)

• constraints (tuple[str, ...])

• identifiability_conditions (tuple[str, ...])

• completeness_basis (str)

• labeled_constraints (int)

• residual_loss_initial (float)

• residual_loss_final (float)

to_dict()

Return a JSON-compatible receipt.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Restore a receipt produced by to_dict().

Parameters:

payload (Any)

Return type:

PINNConstraintReceipt

class QATWrapper(*args, **kwargs)

Bases: object

Torch-absent stand-in: keeps QATWrapper always importable (mixle.models imports it unconditionally, matching every other torch-optional model in this package – see GaussianProcessRegressor), raising only when actually constructed, since it must subclass nn.Module to be real and nn does not exist here.

Parameters:

• args (Any)

• kwargs (Any)

apply_qat(model, *, bits=4, clip_percentile=None, reconfigure=False)

Replace every nn.Linear under model in place with a QATWrapper, so the whole model trains quantization-aware (straight-through int4 fake-quant on every Linear weight, every forward call). Mirrors mixle.experimental.program.lora’s wrapping pattern: walk named_children, swap Linear leaves, recurse into everything else. Weight-tied layers (e.g. CausalLM.head sharing tok.weight) wrap cleanly – only the wrapped module’s forward changes, the underlying nn.Parameter (and anything else pointing at it) is untouched. Returns model (mutated in place) for chaining.

Parameters:

• model (Any)

• bits (int)

• clip_percentile (float | None)

• reconfigure (bool)

Return type:

Any

fake_quantize(x, *, bits=4, clip_percentile=None)

Straight-through fake-quantize x to bits (int4 or int8, per mixle.task.quantize._QMAX): forward returns the real quantize->dequantize round trip, backward passes the gradient through unchanged (STE).

Parameters:

• x (Any)

• bits (int)

• clip_percentile (float | None)

Return type:

Any

fake_quantize_int4(x, *, clip_percentile=None)

fake_quantize(x, bits=4) – the int4 case this roadmap item targets.

Parameters:

• x (Any)

• clip_percentile (float | None)

Return type:

Any

set_fake_quant_enabled(model, enabled)

Toggle every QATWrapper under model on/off in place. enabled=False runs the model at its real fp32 weights (e.g. to check that QAT training didn’t wreck full-precision quality); enabled=True (the default after apply_qat()) restores the fake-quant forward. Returns model for chaining.

Parameters:

• model (Any)

• enabled (bool)

Return type:

Any

class PoissonRegressionNeuralNetwork(module, engine=None, precision=None)

Bases: object

A Torch count-regression wrapper optimized by Poisson log likelihood.

The wrapped module predicts log rates. Observed counts must be non-negative and match the module output shape after one-dimensional inputs are promoted to column vectors.

Parameters:

• module (Any)

• engine (Any | None)

• precision (Any | None)

parameters()

Return trainable parameters of the wrapped log-rate module.

Return type:

Iterable[Any]

log_rate_tensor(x)

Return predicted log rates as a Torch tensor.

Parameters:

x (Any)

Return type:

Any

log_likelihood(x, y)

Return the summed Poisson count log likelihood.

Parameters:

• x (Any)

• y (Any)

Return type:

Any

fit(x, y, max_its=500, lr=0.01, optimizer='adam', tol=1.0e-7, out=None, print_iter=100, return_result=False, restore_best=True)

Maximize the Poisson count log likelihood.

Parameters:

• x (Any)

• y (Any)

• max_its (int)

• lr (float)

• optimizer (str)

• tol (float)

• out (Any | None)

• print_iter (int)

• return_result (bool)

• restore_best (bool)

Return type:

Any

predict_rate_tensor(x)

Return predicted Poisson rates as a Torch tensor.

Parameters:

x (Any)

Return type:

Any

predict_rate(x)

Return predicted Poisson rates as a NumPy array.

Parameters:

x (Any)

Return type:

ndarray

predict(x)

Return rounded count predictions as integer NumPy values.

Parameters:

x (Any)

Return type:

ndarray

class RandomForestConditional(forest, task, sigma=None, n_features=None, name=None, keys=None, forest_spec=None)

Bases: SequenceEncodableProbabilityDistribution

Fitted random forest viewed as a conditional distribution p(y | x).

Observations are (x, y) pairs: x is a feature vector and y is a class label (classification) or a real target (regression). seq_log_density returns log p(y | x) – predict_log_proba for classification, a Gaussian residual density with scale sigma for regression.

Parameters:

• forest (Any)

• task (str)

• sigma (float | None)

• n_features (int | None)

• name (str | None)

• keys (str | None)

• forest_spec (dict[str, Any] | None)

density(x)

Return p(y | x) for one feature/target pair.

Parameters:

x (tuple[Any, Any])

Return type:

float

log_density(x)

Return log p(y | x) for one feature/target pair.

Parameters:

x (tuple[Any, Any])

Return type:

float

seq_log_density(x)

Return per-row conditional log densities for encoded (X, y) data.

Parameters:

x (tuple[ndarray, ndarray])

Return type:

ndarray

sample_y(x, rng)

Draw target values from the fitted conditional forest at feature rows x.

Parameters:

• x (Any)

• rng (RandomState)

Return type:

ndarray

sampler(seed=None)

Return a conditional sampler for drawing targets given features.

Parameters:

seed (int | None)

Return type:

RandomForestConditionalSampler

estimator(pseudo_count=None)

Return an estimator preserving the fitted model’s complete specification.

Parameters:

pseudo_count (float | None)

Return type:

RandomForestEstimator

dist_to_encoder()

Return the encoder for feature/target observation pairs.

Return type:

RandomForestEncoder

class RandomForestEstimator(task=None, n_estimators=100, max_depth=None, min_samples_split=2, min_samples_leaf=1, max_features='auto', random_state=None, min_sigma=1.0e-3, name=None, keys=None, n_features=None)

Bases: ParameterEstimator

Estimator that fits a native (numpy) random forest as a conditional leaf.

task is explicitly ‘classification’ or ‘regression’; target dtype is never used to guess semantics. The forest hyperparameters (n_estimators, max_depth, min_samples_split, min_samples_leaf, max_features, random_state) are passed straight to the native ensemble. estimate() trains in one pass on the accumulated weighted data; there is no EM iteration, so drive it with optimize(max_its=1) or call the seq_encode / accumulate / estimate path directly.

Parameters:

• task (str | None)

• n_estimators (int)

• max_depth (int | None)

• min_samples_split (int)

• min_samples_leaf (int)

• max_features (Any)

• random_state (int | None)

• min_sigma (float)

• name (str | None)

• keys (str | None)

• n_features (int | None)

accumulator_factory()

Return an accumulator factory for weighted feature/target buffers.

Return type:

RandomForestAccumulatorFactory

estimate(nobs, suff_stat)

Fit the native forest from buffered data and return it as a conditional leaf.

Parameters:

• nobs (float | None)

• suff_stat (tuple[ndarray, ndarray, ndarray] | None)

Return type:

RandomForestConditional

class StochasticBlockGraphModel(block_probs, block_assignments, directed=False, self_loops=False, name=None)

Bases: object

Bernoulli stochastic block model with fixed node assignments.

Parameters:

• block_probs (Any)

• block_assignments (Sequence[int])

• directed (bool)

• self_loops (bool)

• name (str | None)

classmethod fit_mle(adjacency, block_assignments, num_blocks=None, directed=False, self_loops=False, pseudo_count=0.0, prior_p=0.5, name=None)

Thin shim delegating to fit_stochastic_block_mle (kept for the classmethod-fit call API).

Parameters:

• adjacency (Any)

• block_assignments (Sequence[int])

• num_blocks (int | None)

• directed (bool)

• self_loops (bool)

• pseudo_count (float)

• prior_p (float)

• name (str | None)

Return type:

StochasticBlockGraphModel

log_likelihood(adjacency)

Return the Bernoulli SBM log likelihood.

Parameters:

adjacency (Any)

Return type:

float

sample(seed=None)

Draw one graph from the block model.

Parameters:

seed (int | None)

Return type:

ndarray

bic(adjacency)

BIC using the number of identifiable block edge probabilities.

Parameters:

adjacency (Any)

Return type:

float

class TransEKnowledgeGraphModel(entity_embeddings, relation_embeddings, entity_names=None, relation_names=None, name=None)

Bases: object

Dependency-free TransE model with a NumPy margin objective.

Parameters:

• entity_embeddings (Any)

• relation_embeddings (Any)

• entity_names (Sequence[Any] | None)

• relation_names (Sequence[Any] | None)

• name (str | None)

classmethod random(num_entities, num_relations, embedding_dim=16, seed=None, scale=0.01, entity_names=None, relation_names=None, name=None)

Create a randomly initialized model.

Parameters:

• num_entities (int)

• num_relations (int)

• embedding_dim (int)

• seed (int | None)

• scale (float)

• entity_names (Sequence[Any] | None)

• relation_names (Sequence[Any] | None)

• name (str | None)

Return type:

TransEKnowledgeGraphModel

distance_triples(triples)

Return squared TransE distances ||h + r - t||^2.

Parameters:

triples (Sequence[tuple[Any, Any, Any]])

Return type:

ndarray

score_triples(triples)

Return TransE scores; higher is more plausible.

Parameters:

triples (Sequence[tuple[Any, Any, Any]])

Return type:

ndarray

margin_loss(positive_triples, negative_triples, margin=1.0)

Return the pairwise TransE ranking loss.

Parameters:

• positive_triples (Sequence[tuple[Any, Any, Any]])

• negative_triples (Sequence[tuple[Any, Any, Any]])

• margin (float)

Return type:

float

negative_sample(triples, seed=None, corrupt='tail', *, known_triples=None)

Produce filtered corruptions that are not any supplied known true triple.

Parameters:

• triples (Sequence[tuple[Any, Any, Any]])

• seed (int | None)

• corrupt (str)

• known_triples (Sequence[tuple[Any, Any, Any]] | None)

Return type:

list[tuple[Any, Any, Any]]

fit_margin(positive_triples, negative_triples=None, margin=1.0, lr=0.01, max_its=100, seed=None, normalize_entities=True)

Fit embeddings with simple stochastic subgradient descent.

Parameters:

• positive_triples (Sequence[tuple[Any, Any, Any]])

• negative_triples (Sequence[tuple[Any, Any, Any]] | None)

• margin (float)

• lr (float)

• max_its (int)

• seed (int | None)

• normalize_entities (bool)

Return type:

KnowledgeGraphFitResult

normalize_entity_embeddings(max_norm=1.0)

Project entity embeddings into an L2 ball.

Parameters:

max_norm (float)

Return type:

None

class TruncatedDirichletProcessMixtureFitResult(model, responsibilities, history, objective_name='truncated_dp_elbo')

Bases: object

Fitted truncated DPM plus variational responsibilities and history.

Parameters:

• model (TruncatedDirichletProcessMixtureModel)

• responsibilities (ndarray)

• history (list[float])

• objective_name (str)

class TruncatedDirichletProcessMixtureModel(components, alpha=1.0, gamma=None, weights=None, name=None)

Bases: object

Truncated stick-breaking mixture over existing mixle component models.

Parameters:

• components (Sequence[SequenceEncodableProbabilityDistribution])

• alpha (float)

• gamma (Any | None)

• weights (Any | None)

• name (str | None)

property expected_log_weights: ndarray

Return E_q[log pi_k] under the variational stick posteriors.

component_log_density(x)

Return component log densities for one observation.

Parameters:

x (Any)

Return type:

ndarray

log_density(x)

Return the finite-truncation mixture log density for one observation.

Parameters:

x (Any)

Return type:

float

density(x)

Return the finite-truncation mixture density for one observation.

Parameters:

x (Any)

Return type:

float

responsibilities(data, expected=True)

Return posterior component probabilities for observations.

Parameters:

• data (Sequence[Any])

• expected (bool)

Return type:

ndarray

effective_components(threshold=0.01)

Count components with posterior mean stick weight above threshold.

Parameters:

threshold (float)

Return type:

int

sample(size=None, seed=None)

Draw observations from the finite truncation.

Parameters:

• size (int | None)

• seed (int | None)

Return type:

Any | list[Any]

baum_welch_pomdp(sequences, num_states, num_actions, num_observations, initial_model=None, max_its=50, tol=1.0e-8, pseudo_count=1.0e-3, seed=None)

Fit a known-action finite PartiallyObservableMarkovDecisionProcess by Baum-Welch/EM.

Parameters:

• sequences (Sequence[tuple[Sequence[int], Sequence[int]]])

• num_states (int)

• num_actions (int)

• num_observations (int)

• initial_model (PartiallyObservableMarkovDecisionProcessModel | None)

• max_its (int)

• tol (float | None)

• pseudo_count (float)

• seed (int | None)

Return type:

PartiallyObservableMarkovDecisionProcessFitResult

discrete_conditional_mutual_information(data, x, y, given=())

Estimate I(X;Y | Z) from categorical samples using empirical counts.

Parameters:

• data (Any)

• x (int)

• y (int)

• given (Sequence[int])

Return type:

float

expected_log_stick_weights(gamma)

Return E_q[log pi_k] for truncated Beta stick posteriors.

Parameters:

gamma (Any)

Return type:

ndarray

fit_erdos_renyi_mle(adjacency, directed=False, self_loops=False, pseudo_count=0.0, prior_p=0.5, name=None)

Conjugate-Bernoulli MLE of the edge probability (module-level estimation, not a classmethod-fit).

Parameters:

• adjacency (Any)

• directed (bool)

• self_loops (bool)

• pseudo_count (float)

• prior_p (float)

• name (str | None)

Return type:

ErdosRenyiGraphModel

fit_stochastic_block_mle(adjacency, block_assignments, num_blocks=None, directed=False, self_loops=False, pseudo_count=0.0, prior_p=0.5, name=None)

Conjugate-Bernoulli MLE of block edge probabilities for fixed assignments (module-level estimation).

Parameters:

• adjacency (Any)

• block_assignments (Sequence[int])

• num_blocks (int | None)

• directed (bool)

• self_loops (bool)

• pseudo_count (float)

• prior_p (float)

• name (str | None)

Return type:

StochasticBlockGraphModel

fit_truncated_dpm(data, initial_components, component_estimator, alpha=1.0, max_its=50, tol=1.0e-8, sort_components=True, name=None)

Fit a truncated DP mixture by coordinate-ascent variational updates.

The component M-steps are delegated to ordinary mixle.stats estimators. Components are point-estimated parameters; history is the complete truncated-stick ELBO over assignments and Beta sticks conditional on those point estimates (expected likelihood + assignment entropy - stick KL).

Parameters:

• data (Sequence[Any])

• initial_components (Sequence[SequenceEncodableProbabilityDistribution])

• component_estimator (ParameterEstimator | Sequence[ParameterEstimator])

• alpha (float)

• max_its (int)

• tol (float | None)

• sort_components (bool)

• name (str | None)

Return type:

TruncatedDirichletProcessMixtureFitResult

fit_induced_pcfg(data, terminal_estimators, max_nonterminals, initial_model=None, vdata=None, max_its=10, init_p=1.0, seed=None, terminal_rule_mass=0.5, rule_pseudo_count=1.0e-3, prune_threshold=0.0, min_rule_prob=0.0, start='S', name=None)

Fit an induced heterogeneous PCFG and track train/validation likelihoods.

Parameters:

• data (Sequence[Sequence[Any]])

• terminal_estimators (Sequence[ParameterEstimator])

• max_nonterminals (int)

• initial_model (HeterogeneousPCFGDistribution | None)

• vdata (Sequence[Sequence[Any]] | None)

• max_its (int)

• init_p (float)

• seed (int | None)

• terminal_rule_mass (float)

• rule_pseudo_count (float | None)

• prune_threshold (float)

• min_rule_prob (float)

• start (Any)

• name (str | None)

Return type:

GrammarLearningResult

gaussian_conditional_independence(data, x, y, given=(), alpha=0.05, ridge=1.0e-10)

Fisher-z Gaussian conditional independence test.

Parameters:

• data (Any)

• x (int)

• y (int)

• given (Sequence[int])

• alpha (float)

• ridge (float)

Return type:

ConditionalIndependenceResult

gaussian_partial_correlation(data, x, y, given=(), ridge=1.0e-10)

Return partial correlation rho_xy.given for continuous data.

Parameters:

• data (Any)

• x (int)

• y (int)

• given (Sequence[int])

• ridge (float)

Return type:

float

grammar_rule_table(model)

Return a flat, inspectable rule table for learned PCFGs.

Parameters:

model (HeterogeneousPCFGDistribution)

Return type:

list[dict[str, Any]]

hard_em_stochastic_block_model(adjacency, num_blocks, max_its=20, restarts=1, seed=None, directed=False, self_loops=False, pseudo_count=1.0, prior_p=0.5)

Classification/hard-EM fit for a stochastic block model.

Parameters:

• adjacency (Any)

• num_blocks (int)

• max_its (int)

• restarts (int)

• seed (int | None)

• directed (bool)

• self_loops (bool)

• pseudo_count (float)

• prior_p (float)

Return type:

HardEMResult

leapfrog_rollout(net, q0, p0, dt, n_steps)

Leapfrog integration of net’s learned (or untrained) Hamiltonian flow.

Returns (qs, ps), each shaped (n_steps + 1, *q0.shape) – the trajectory including the initial state. This scheme is available only when net.is_separable is true. A symplectic integrator is the right numerical counterpart to a conservative continuous-time system: unlike a generic (e.g. Euler) integrator, its energy error stays bounded and oscillates rather than drifting monotonically over long rollouts.

Parameters:

• net (HamiltonianNet)

• q0 (Any)

• p0 (Any)

• dt (float)

• n_steps (int)

Return type:

tuple[Any, Any]

learn_pc_skeleton(data, variable_names=None, alpha=0.05, max_cond_set=2, method='gaussian')

Learn a PC-style undirected skeleton from conditional independences.

Parameters:

• data (Any)

• variable_names (Sequence[Any] | None)

• alpha (float)

• max_cond_set (int)

• method (str)

Return type:

CausalSkeleton

build_autoregressive_categorical(dim, n_categories, *, hidden=64)

An autoregressive neural density over discrete vectors x in {0..C-1}^dim – exact, normalized p(x).

The continuous flows/VAE above model R^d; heterogeneous data is also categorical. This factorizes p(x) = prod_i p(x_i | x_{<i}) with a MADE-masked network whose per-coordinate softmax is each conditional, so the density is exactly normalized (sums to 1 over the finite space) and can be compared directly with count/categorical families. log_density sums the picked log-softmax logits; sample fills the vector one coordinate at a time. Another ready module for NeuralDensity; the adapter is unchanged.

Parameters:

• dim (int)

• n_categories (int)

• hidden (int)

Return type:

Any

build_contrastive_projection(d_x, d_y, *, encoder_x=None, encoder_y=None, proj_dim=None, hidden=64, freeze_encoders=True, temperature=0.07)

Build a contrastive projection/ranker between two embedding spaces.

This is the stage-1 multimodal pattern – frozen encoder -> trainable projection -> frozen encoder – stated with no domain nouns. encoder_x/encoder_y are any torch module mapping a raw item to a d_x/d_y embedding; both default to nn.Identity(), so x/y may already BE the embeddings (pass precomputed vectors straight in, no backbone required). Encoders are frozen by default (freeze_encoders=True): their parameters get requires_grad_(False) and the module is pinned in eval() regardless of the outer train()/eval() calls the M-step makes, so no dropout/batchnorm noise leaks into a “frozen” backbone and no gradient ever reaches it. The only trainable piece is a small projection head per side (d_x / d_y -> hidden -> proj_dim, default proj_dim = min(d_x, d_y)) mapping BOTH embeddings into one shared, L2-normalized space – the CLIP design (two projections into a shared space), not a single asymmetric x -> y regression – so the same leaf answers “which y matches this x” and “which x matches this y”.

contrastive_loss is a symmetric weighted InfoNCE objective over a complete paired batch. Anchor weights and the two candidate/negative measures are explicit and independently validated, so responsibilities are never applied only after cross-row negative contributions have already been formed. score_pairs and similarity_matrix are batch-independent retrieval surfaces.

This object intentionally has no log_density or sample_given method: a contrastive batch objective is not a normalized per-observation probability distribution and cannot be composed into mixture/HMM EM.

Parameters:

• d_x (int)

• d_y (int)

• encoder_x (Any)

• encoder_y (Any)

• proj_dim (int | None)

• hidden (int)

• freeze_encoders (bool)

• temperature (float)

Return type:

Any

build_conditional_autoregressive_categorical(x_dim, y_dim, n_categories, *, hidden=64)

An autoregressive categorical conditioned on x: exact p(y | x) over discrete y in {0..C-1}^y_dim.

The conditional sibling of build_autoregressive_categorical() and the discrete counterpart to build_conditional_flow(). It factorizes p(y | x) = prod_i p(y_i | y_{<i}, x) with a MADE-masked net over y into which x is injected unmasked (degree 0, so every coordinate may depend on x). Each per-coordinate softmax is exactly a conditional, so the density is exactly normalized and comparable to other exact discrete conditional leaves. Exposes log_density(x, y) and sample_given(x) – the contract a NeuralConditionalDensity adapts.

Parameters:

• x_dim (int)

• y_dim (int)

• n_categories (int)

• hidden (int)

Return type:

Any

build_conditional_flow(x_dim, y_dim, *, hidden=32, layers=4)

A conditional coupling flow: an exact p(y | x) whose transform of y is conditioned on x.

The exact-density counterpart to build_mdn(). Each affine-coupling layer’s shift/scale networks take both the passed-through y coordinates and x, so the whole invertible y-transform bends with the input – capturing within-``y`` dependence (e.g. y2 a nonlinear function of y1) that a single-Gaussian NeuralGaussian (isotropic mean-only) cannot, while keeping an exact log-density rather than a bound. Needs y_dim >= 2 for the coupling to be non-trivial. Exposes log_density(x, y) and sample_given(x) – the contract a NeuralConditionalDensity adapts.

Parameters:

• x_dim (int)

• y_dim (int)

• hidden (int)

• layers (int)

Return type:

Any

build_convex_energy_net(dim, *, hidden=64, layers=3)

An input-convex MLP energy E(x): R^dim -> R, convex in x BY CONSTRUCTION – ready to wrap in an EnergyModel exactly like build_energy_net(). A convex energy gives Langevin sampling (EnergyModelSampler) a unimodal target with no spurious local minima to get stuck in, and gives any consumer of the fitted energy a certified-convex scalar-valued potential (e.g. a verified optimum for a downstream mixle.doe search over -E(x)). See _convex_energy_net_class() for the construction.

Parameters:

• dim (int)

• hidden (int)

• layers (int)

Return type:

Any

build_coupling_flow(dim, *, hidden=32, layers=4)

A RealNVP coupling flow over R^dim with an exact log_density(x) and sample(n) – ready to wrap.

Alternating affine-coupling layers map data to a standard-normal base; log_density is the base log-prob plus the log-determinant of the (triangular) Jacobian. A minimal, correct instance of the density module a NeuralDensity adapts – swap in any other module with the same two methods.

Parameters:

• dim (int)

• hidden (int)

• layers (int)

Return type:

Any

build_energy_net(dim, *, hidden=64, layers=3)

An MLP energy E(x): R^dim -> R plus a learned scalar normalization offset.

Lower energy = higher unnormalized density. log_norm uses the -log Z convention, so the paired EnergyModel scores -E(x) + log_norm. Swap in any module exposing energy(x) -> (n,), a scalar log_norm parameter, and a dim attribute.

Parameters:

• dim (int)

• hidden (int)

• layers (int)

Return type:

Any

build_maf(dim, *, hidden=64, blocks=3)

A masked autoregressive flow over R^dim – an exact density that factorizes p(x) by the chain rule, each p(x_i | x_{<i}) an affine map with autoregressive (MADE-masked) mean and log-scale.

Unlike the coupling flow it conditions every coordinate on all earlier ones (a richer autoregressive dependence), and unlike the VAE its log_density is exact. It can therefore be compared directly with a Gaussian, a flow, or another exact-density leaf. Sampling is the sequential inverse (one coordinate at a time). Another ready module for NeuralDensity; the adapter is unchanged.

Parameters:

• dim (int)

• hidden (int)

• blocks (int)

Return type:

Any

build_mdn(x_dim, y_dim, *, k=5, hidden=32, layers=2)

A mixture density network: p(y | x) = sum_k pi_k(x) N(y; mu_k(x), diag sigma_k(x)^2) – ready to wrap.

A shared MLP body maps x to three heads – mixing logits, component means, and (log) component scales – so the entire conditional law is a function of x: multimodal (several mu_k) and heteroscedastic (input-dependent sigma_k). Exposes log_density(x, y) (a log-sum-exp over components) and sample_given(x) (pick a component by pi, then a Gaussian), the contract a NeuralConditionalDensity adapts.

Parameters:

• x_dim (int)

• y_dim (int)

• k (int)

• hidden (int)

• layers (int)

Return type:

Any

build_product_energy_net(experts)

Combine expert energy modules multiplicatively: one module with energy(x) = sum_k experts[k].energy(x).

A product of experts, p(x) ∝ prod_k p_k(x) – a conjunction (each expert a soft constraint, the product their intersection), as opposed to a mixture’s disjunction. This is the ENERGY-BASED, arbitrary-density complement to mixle.ops.product_of_experts(), which pools tractable families (Categorical, Gaussian) in closed form but deliberately raises on the general continuous case because the product normalizer is then intractable. That intractable normalizer is exactly what the energy stack already handles: wrap the result here in an EnergyModel to fit the shared log_norm by NCE and sample by Langevin, no new machinery.

Each expert must expose energy(x) -> (n,) and a dim attribute (e.g. any build_energy_net() / build_convex_energy_net() module, all sharing one input dim), and stays individually inspectable via the built module’s .experts / .expert_energies(x). Fit it in one line:

model = EnergyModel(build_product_energy_net([expert_a, expert_b]), m_steps=250)

Parameters:

experts (Any)

Return type:

Any

build_projection_leaf(d_x, d_y, *, encoder_x=None, encoder_y=None, proj_dim=None, hidden=64, freeze_encoders=True, temperature=0.07)

Compatibility spelling for build_contrastive_projection(); the result is not a distribution leaf.

Parameters:

• d_x (int)

• d_y (int)

• encoder_x (Any)

• encoder_y (Any)

• proj_dim (int | None)

• hidden (int)

• freeze_encoders (bool)

• temperature (float)

Return type:

Any

build_vae(dim, *, latent=2, hidden=32, eval_samples=16)

Build a variational autoencoder over R^dim.

An amortized encoder q(z | x) and a decoder p(x | z) (diagonal-Gaussian, learned observation scale) are trained by the ELBO with the reparameterization trick. This is a different family from a flow: structure is represented through a low-dimensional latent rather than an invertible map, while the same NeuralDensity adapter can still use it because it exposes the same two methods.

At evaluation, log_density(x) returns a deterministic Monte Carlo estimate of the ELBO using eval_samples fixed local draws from q(z | x). elbo_estimate(x) returns that estimate together with its Monte Carlo standard error. It is an estimate of a lower bound, not an exact log density and not itself guaranteed to remain below log p(x) at finite sample count. Compare VAE leaves with other bounded leaves whenever possible. Training uses fresh reparameterized samples for an unbiased gradient.

Parameters:

• dim (int)

• latent (int)

• hidden (int)

• eval_samples (int)

Return type:

Any

make_deep_set(element_dim, phi_hidden, latent_dim, rho_hidden, output_dim=1, *, pooling='mean')

A Deep Sets network (Zaheer et al. 2017): invariant to any permutation of the set axis, by construction.

Input shape (..., set_size, element_dim): a per-element MLP phi (shared weights, applied identically to every element – torch.nn.Linear already broadcasts over all leading dims, so reusing make_mlp() for phi gives exactly that) maps each element to a latent_dim code; a permutation-invariant pool (pooling="mean"/"sum"/"max", taken over the set axis) aggregates the codes into one order-independent summary; a second MLP rho maps the summary to the output. Because phi is applied identically per element and the pool is a symmetric function, the output is exactly unchanged by any permutation of the set axis – true for any weights, trained or not, unlike e.g. training on many random orderings and hoping the network learns invariance.

The returned module is a plain torch.nn.Module, trainable with any ordinary Torch optimizer loop over (set_size, element_dim)-shaped inputs. Note: NeuralGaussian’s accumulator flattens each observation to a 1-D feature vector (reshape(n, -1)) before the M-step, which destroys the set axis this module needs – so it is not a drop-in wrapper for set-shaped data as make_mlp()/make_monotonic_mlp() are for flat feature vectors. Use this module directly with a custom training loop (or through a wrapper that preserves the set axis) for a fixed set size.

Parameters:

• element_dim (int)

• phi_hidden (Sequence[int])

• latent_dim (int)

• rho_hidden (Sequence[int])

• output_dim (int)

• pooling (str)

Return type:

Any

make_mlp(input_dim, hidden_dims, output_dim=1, activation='tanh')

Create a simple fully connected Torch MLP.

Parameters:

• input_dim (int)

• hidden_dims (Sequence[int])

• output_dim (int)

• activation (str)

Return type:

Any

make_monotonic_mlp(input_dim, hidden_dims, output_dim=1, *, increasing=True)

A fully connected Torch MLP that is monotonic in every input dimension jointly, BY CONSTRUCTION.

Each layer’s weight matrix is reparameterized through softplus before use, so every weight is strictly non-negative; composed with the (smooth, strictly increasing) Softplus activation, a non-negative-weight affine map followed by an increasing activation is itself increasing, and that property is closed under composition – so the whole network is provably non-decreasing in every input coordinate, with no penalty term and no post-hoc check needed. increasing=False negates the output, giving a network non-increasing in every coordinate instead.

This is a hard architectural constraint (unlike PINNRegression’s soft residual penalty): the guarantee holds at every point in input space, not just where training data landed. Drops into the same wrappers as make_mlp() – NeuralGaussian for regression, NeuralCategorical for classification – no other changes needed. Only jointly monotonic in ALL inputs; a network monotonic in some coordinates and free in others needs a two-path (monotonic + unconstrained) variant, not built here.

Parameters:

• input_dim (int)

• hidden_dims (Sequence[int])

• output_dim (int)

• increasing (bool)

Return type:

Any

mean_stick_weights(gamma)

Return E_q[pi_k] under independent Beta stick posteriors.

Parameters:

gamma (Any)

Return type:

ndarray

orient_v_structures(skeleton)

Build a deterministic, conflict-checked CPDAG from skeleton and separating evidence.

Parameters:

skeleton (CausalSkeleton)

Return type:

PartiallyDirectedGraph

pcfg_log_likelihood(model, data)

Return total PCFG log likelihood on raw sequences.

Parameters:

• model (HeterogeneousPCFGDistribution)

• data (Sequence[Sequence[Any]])

Return type:

float

sample_crp_assignments(num_obs, alpha, seed=None)

Sample Chinese-restaurant-process assignments and table counts.

Parameters:

• num_obs (int)

• alpha (float)

• seed (int | None)

Return type:

tuple[ndarray, ndarray]

stick_breaking_weights(stick_fractions, residual=True)

Convert stick fractions into mixture weights.

When residual is true, the returned vector has one extra final entry containing the remaining stick mass. This is the usual finite truncation.

Parameters:

• stick_fractions (Any)

• residual (bool)

Return type:

ndarray

viterbi_parse(model, sequence)

Return the maximum-probability CKY parse under a heterogeneous PCFG.

Parameters:

• model (HeterogeneousPCFGDistribution)

• sequence (Sequence[Any])

Return type:

PCFGParseNode

Submodules

• mixle.models.coarsening module
  • The implemented model moves and representation helper
  • H1 is this operator inverted
• mixle.models.compress module
  • Three method FAMILIES, one entry point
  • Build vs. borrow
• mixle.models.continual module
• mixle.models.dependence module
• mixle.models.dirichlet_process_mixture module
• mixle.models.dpo_leaf module
• mixle.models.embedding module
• mixle.models.energy module
• mixle.models.eval_harness module
• mixle.models.feature_map module
• mixle.models.gaussian_process module
• mixle.models.grad_leaf module
• mixle.models.grammar module
• mixle.models.hamiltonian module
• mixle.models.knowledge_graph module
• mixle.models.language_model module
• mixle.models.memory_efficient_training module
• mixle.models.mixture_density module
• mixle.models.moe module
• mixle.models.moment_propagation module
  • What is exact vs. approximate
  • Execution and retained-memory contract
• mixle.models.mup module
• mixle.models.neural module
• mixle.models.neural_density module
• mixle.models.neural_families module
• mixle.models.neural_leaf module
• mixle.models.optimizer_routing module
• mixle.models.partially_observable_markov_decision_process module
• mixle.models.pinn module
• mixle.models.qat module
• mixle.models.quotient module
• mixle.models.random_forest module
• mixle.models.random_graph module
• mixle.models.self_distillation module
  • The idea, concretely
  • Why this belongs at the loss-hook level, not a new trainer
• mixle.models.sigma_weighted_projection module
• mixle.models.softmax_leaf module
• mixle.models.sorted_profile_quantizer module
• mixle.models.sparse_gaussian_process module
• mixle.models.sparsity_2_4 module
• mixle.models.streaming_transformer_leaf module
• mixle.models.train_search module
• mixle.models.transformer module
• mixle.models.unified_quantizer module