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)[source]

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()[source]

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)[source]

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

Parameters:

payload (dict)

Return type:

LM

save(path)[source]

Persist the trained LM to path via torch.save (hyperparameters + weights).

Parameters:

path (str)

Return type:

None

classmethod load(path)[source]

Load an LM previously written by save().

Parameters:

path (str)

Return type:

LM

fit(token_ids, *, epochs=1, batch_size=64, lr=3e-3, distributed=False, precision='fp32', shuffle=True, tp_size=1, pp_size=1, cp_size=1)[source]

Pretrain (or continue) on a token-id array via the streaming estimator; the corpus is never buffered.

tp_size/pp_size/cp_size (only meaningful with distributed=True) are the F1 N-D parallelism dimensions – tensor/pipeline/context parallel, ORTHOGONAL to the data-parallel axis (DDP on CPU, FSDP2/ZeRO-3 on CUDA) this handle already runs. They default to 1 (off): the plan is validated against the model’s real dimensions (n_head % tp_size, pp_size <= n_layer, block % cp_size) so a bad plan fails fast, using the sharding/reconstruction mechanism in mixle.utils.parallel.tensor_pipeline_context_parallel (tested there at small scale against the dense forward). Composing the validated plan into real per-axis multi-GPU process groups is the piece that needs actual hardware this environment does not have – see that module’s docstring.

Parameters:
Return type:

LM

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

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:
Return type:

LM

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

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:
Return type:

list

nll(token_ids)[source]

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

Parameters:

token_ids (Any)

Return type:

float

class CategoricalEmbedding(num_categories, dim, *, name=None)[source]

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)

module()[source]

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

Return type:

Any

class CausalSkeleton(edges, separating_sets, variable_names)[source]

Bases: object

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

Parameters:
has_edge(i, j)[source]

Return whether the undirected skeleton contains edge i–j.

Parameters:
Return type:

bool

class CategoricalClassificationNeuralNetwork(module, engine=None, precision=None)[source]

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()[source]

Return trainable parameters of the wrapped classification module.

Return type:

Iterable[Any]

logits_tensor(x)[source]

Return raw class logits for x as a Torch tensor.

Parameters:

x (Any)

Return type:

Any

log_likelihood(x, y)[source]

Return the summed categorical log likelihood for integer labels.

Parameters:
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)[source]

Maximize the categorical classification log likelihood.

Parameters:
Return type:

Any

predict_proba_tensor(x)[source]

Return class probabilities for x as a Torch tensor.

Parameters:

x (Any)

Return type:

Any

predict_proba(x)[source]

Return class probabilities for x as a NumPy array.

Parameters:

x (Any)

Return type:

ndarray

predict(x)[source]

Return maximum-probability class labels for x.

Parameters:

x (Any)

Return type:

ndarray

class ConditionalIndependenceResult(measure, statistic, p_value, independent)[source]

Bases: object

Result from a conditional independence calculation.

Parameters:
DPOLeaf

alias of DPOModel

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

Bases: object

Independent Bernoulli edge model for directed or undirected graphs.

Parameters:
classmethod fit_mle(adjacency, directed=False, self_loops=False, pseudo_count=0.0, prior_p=0.5, name=None)[source]

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

Parameters:
Return type:

ErdosRenyiGraphModel

log_likelihood(adjacency)[source]

Return the Bernoulli graph log likelihood.

Parameters:

adjacency (Any)

Return type:

float

sample(num_nodes, seed=None)[source]

Draw one binary adjacency matrix.

Parameters:

  • num_nodes (int)

  • seed (int | None)

Return type:

ndarray

bic(adjacency)[source]

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)[source]

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:
parameters()[source]

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)[source]

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

Parameters:
Return type:

Any

log_marginal_likelihood(x, y)[source]

Return the exact GP log marginal likelihood for training data.

Parameters:
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)[source]

Maximize the GP log marginal likelihood.

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

Parameters:
Return type:

Any

predict(x_train, y_train, x_new, return_cov=False)[source]

Return posterior predictive mean, and optionally covariance.

Parameters:
Return type:

Any

predict_monotone(x_train, y_train, x_new, increasing=True)[source]

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:
Return type:

ndarray

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

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()[source]

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)[source]

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

Parameters:

x (Any)

Return type:

Any

log_likelihood(x, y)[source]

Return the summed Gaussian regression log likelihood.

Parameters:
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)[source]

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:
Return type:

Any

predict(x)[source]

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')[source]

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; device (e.g. "mps"/"cuda") runs it on the GPU.

Parameters:

  • module (Any)

  • m_steps (int)

  • lr (float)

  • name (str | None)

  • batch_size (int | None)

  • device (str)

log_density(xy)[source]

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

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)[source]

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

Parameters:

enc (Any)

Return type:

ndarray

predict(x)[source]

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

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return a conditional sampler over labels given features.

Parameters:

seed (int | None)

Return type:

NeuralCategoricalSampler

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

NeuralCategoricalEstimator

dist_to_encoder()[source]

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

Return type:

NeuralCategoricalEncoder

to_dict()[source]

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

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

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)[source]

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)[source]

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

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)[source]

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

Parameters:

enc (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return a conditional sampler for drawing y given x.

Parameters:

seed (int | None)

Return type:

NeuralConditionalDensitySampler

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

NeuralConditionalDensityEstimator

dist_to_encoder()[source]

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

Return type:

NeuralConditionalDensityEncoder

to_dict()[source]

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

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

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, precision='fp32', name=None, loss=None, optimizer=None)[source]

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)

  • precision (str)

  • name (str | None)

  • loss (Any)

  • optimizer (Any)

accumulator_factory()[source]

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

Return type:

DataBufferAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a distribution from accumulated sufficient statistics.

Parameters:
Return type:

GradLeaf

class GradLeaf(module, *, m_steps=60, lr=5e-3, device=None, batch_size=None, precision='fp32', name=None, loss=None, optimizer=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Wrap a torch density module (module.log_density(x) -> (n,)) as a composable mixle 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)

  • precision (str)

  • name (str | None)

  • loss (Any)

  • optimizer (Any)

log_density(x)[source]

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

Parameters:

x (Any)

Return type:

float

seq_log_density(x)[source]

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

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

GradLeafSampler

estimator(pseudo_count=None)[source]

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

GradEstimator

dist_to_encoder()[source]

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', name=None)[source]

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)

  • name (str | None)

log_density(x)[source]

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

Parameters:

x (Any)

Return type:

float

seq_log_density(x)[source]

Return per-row log densities for encoded observations.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

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

Parameters:

seed (int | None)

Return type:

NeuralDensitySampler

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

NeuralDensityEstimator

dist_to_encoder()[source]

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

Return type:

NeuralDensityEncoder

to_dict()[source]

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

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

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', name=None)[source]

Bases: ParameterEstimator

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)

  • name (str | None)

accumulator_factory()[source]

Return an accumulator factory for weighted neural-density batches.

Return type:

NeuralDensityAccumulatorFactory

estimate(nobs, suff_stat)[source]

Run the weighted neural-density M-step and return the updated leaf.

Parameters:
Return type:

NeuralDensity

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

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)[source]

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

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)[source]

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

Parameters:

enc (Any)

Return type:

ndarray

classmethod compute_capabilities()[source]

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

backend_seq_log_density(enc, engine)[source]

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

Parameters:
Return type:

Any

sampler(seed=None)[source]

Return a conditional sampler for drawing y given x.

Parameters:

seed (int | None)

Return type:

NeuralGaussianSampler

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

NeuralGaussianEstimator

dist_to_encoder()[source]

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

Return type:

NeuralGaussianEncoder

to_dict()[source]

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

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Rebuild a NeuralGaussian from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralGaussian

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

Bases: _NeuralFamily

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

log_density is the ELBO – a lower bound on log p(x) – evaluated deterministically at the encoder mean so an EM log-likelihood stays monotone. Compare it with other bounded leaves whenever possible; mixing it with an exact-density leaf, such as a Gaussian or flow, compares a bound against an exact value and can under-weight the VAE. See build_vae() for the full statement.

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

Bases: _NeuralFamily

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

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

Bases: _NeuralFamily

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

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

Bases: _NeuralFamily

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

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

Bases: SequenceEncodableProbabilityDistribution

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

Parameters:
seq_log_density(enc)[source]

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

Parameters:

enc (Any)

Return type:

ndarray

log_density(xcr)[source]

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

Parameters:

xcr (Any)

Return type:

float

prefers(x)[source]

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

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return the sampler for the preference-scoring leaf.

Parameters:

seed (int | None)

Return type:

DPOModelSampler

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

DPOModelEstimator

dist_to_encoder()[source]

Return the encoder for preference triples.

Return type:

DPOEncoder

to_dict()[source]

Serialize policy/reference modules and DPO hyperparameters.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

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)[source]

Bases: SequenceEncodableProbabilityDistribution

log p(x) ≈ -E(x) + c for an energy module (module.energy(x) -> (n,) and a learned scalar log_norm).

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)[source]

Return the approximate normalized log density for one observation.

Parameters:

x (Any)

Return type:

float

seq_log_density(x)[source]

Return approximate normalized log densities for a batch of observations.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return a Langevin sampler for the learned energy model.

Parameters:

seed (int | None)

Return type:

EnergyModelSampler

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

EnergyModelEstimator

dist_to_encoder()[source]

Return the encoder for vectorized energy-model scoring.

Return type:

EnergyModelEncoder

to_dict()[source]

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

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Rebuild an EnergyModel from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

EnergyModel

class FeatureMapDensity(feature_name, inner, name=None)[source]

Bases: SequenceEncodableProbabilityDistribution

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

Parameters:

  • feature_name (str)

  • inner (Any)

  • name (str | None)

density(x)[source]

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

Parameters:

x (Any)

Return type:

float

log_density(x)[source]

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

Parameters:

x (Any)

Return type:

float

seq_log_density(x)[source]

Return inner log densities for an already-featurized batch.

Parameters:

x (ndarray)

Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for the inner feature-space distribution.

Parameters:

seed (int | None)

Return type:

FeatureMapSampler

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

FeatureMapEstimator

dist_to_encoder()[source]

Return the encoder that maps raw items to feature vectors.

Return type:

FeatureMapEncoder

class FeatureMapEstimator(feature_name, inner, name=None)[source]

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)

accumulator_factory()[source]

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

Return type:

FeatureMapAccumulatorFactory

estimate(nobs, suff_stat)[source]

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

Parameters:
Return type:

FeatureMapDensity

feature_fn(name)[source]

Look up a registered feature function by name; raises if it was never registered.

Parameters:

name (str)

Return type:

Callable[[Any], ndarray]

register_feature_fn(name, fn)[source]

Register fn (raw item -> fixed-length vector) under name so a leaf can carry just the name.

Parameters:
Return type:

None

class StreamingTransformer(module, device='cpu')[source]

Bases: SequenceEncodableProbabilityDistribution

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

Parameters:

  • module (Any)

  • device (str)

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

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

embedding optionally ties a shared CategoricalEmbedding across leaves.

Parameters:
Return type:

StreamingTransformer

log_density(xy)[source]

Return the next-token log probability for one (context, token) pair.

Parameters:

xy (Any)

Return type:

float

predict(x)[source]

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

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return the sampler for the conditional next-token model.

Parameters:

seed (int | None)

Return type:

StreamingTransformerSampler

seq_log_density(enc)[source]

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

Parameters:

enc (Any)

Return type:

ndarray

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

StreamingTransformerEstimator

dist_to_encoder()[source]

Return the encoder for context/token training pairs.

Return type:

StreamingTokenEncoder

to_dict()[source]

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

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

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)[source]

Bases: object

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

Parameters:
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))[source]

Bases: object

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

Parameters:
dims()[source]

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

Return type:

int

bounds()[source]

Return unit-cube bounds for the DOE optimizer.

Return type:

list[tuple[float, float]]

decode(point)[source]

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')[source]

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:
build_causal_lm(vocab, d_model=128, n_layer=3, n_head=4, block=64, embedding=None, gradient_checkpointing=False)[source]

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

ewc(anchor, fisher, lam=1.0)[source]

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

Parameters:
Return type:

tuple

extrapolate_learning_curve(steps, losses, *, at)[source]

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. Falls back to the last observed loss if the fit fails.

Parameters:
Return type:

float

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

Diagonal empirical Fisher of a classification leaf’s module on (x, y): mean of (d log p(y|x)/dtheta)^2.

Parameters:
Return type:

list

lm_train_fn(token_ids, val_ids, *, vocab, block=64, max_epochs=3, device='cpu')[source]

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

budget in (0, 1] scales the number of epochs. A larger pretraining loop can use the same convention to scale steps or token subsets.

Parameters:
Return type:

Callable[[dict[str, Any], float], float]

snapshot(leaf_or_module)[source]

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

Parameters:

leaf_or_module (Any)

Return type:

list

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

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.

Parameters:
Return type:

TrainingSearchResult

apply_mup_init(model, *, base_width, base_std=0.02)[source]

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:
Return type:

None

classify_causal_lm_params(model)[source]

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)[source]

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)[source]

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)[source]

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:
Return type:

list[dict]

output_forward_multiplier(width_mult)[source]

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')[source]

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')[source]

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)[source]

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

Parameters:
Return type:

tuple

class GrammarLearningResult(model, history, validation_history=None)[source]

Bases: FitResult[HeterogeneousPCFGDistribution]

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

Parameters:
class HamiltonianNet(dim, hidden=(64, 64))[source]

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 MLP computing H from the concatenated (q, p). 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])

hamiltonian(q, p)[source]

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

Parameters:
Return type:

Any

time_derivative(q, p)[source]

(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:
Return type:

tuple[Any, Any]

class HardEMResult(model, history, validation_history=None)[source]

Bases: FitResult[StochasticBlockGraphModel]

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

Parameters:
class KnowledgeGraphFitResult(model, history, validation_history=None)[source]

Bases: FitResult[TransEKnowledgeGraphModel]

Result from TransE margin fitting.

Parameters:
class PartiallyObservableMarkovDecisionProcessFilterResult(beliefs, log_likelihood, predictive_observation_probs)[source]

Bases: object

Belief trajectories, log likelihood, and predictive observation terms.

Parameters:
class PartiallyObservableMarkovDecisionProcessFitResult(model, history, validation_history=None)[source]

Bases: FitResult[PartiallyObservableMarkovDecisionProcessModel]

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

Parameters:
class PartiallyObservableMarkovDecisionProcessModel(transition, observation, initial_belief=None, rewards=None, name=None)[source]

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)[source]

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)[source]

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

Parameters:
Return type:

PartiallyObservableMarkovDecisionProcessFilterResult

sequence_log_likelihood(actions, observations, initial_belief=None)[source]

Return log P(observations | actions, model).

Parameters:
Return type:

float

forward_backward(actions, observations, initial_belief=None)[source]

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

Parameters:
Return type:

tuple[ndarray, ndarray, float]

predict_observation(belief, action)[source]

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

Parameters:
Return type:

ndarray

expected_reward(belief, action)[source]

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

Parameters:
Return type:

float

sample(actions, seed=None, initial_belief=None)[source]

Sample latent states and observations for a fixed action sequence.

Parameters:
Return type:

tuple[ndarray, ndarray]

class PartiallyDirectedGraph(directed_edges, undirected_edges, variable_names)[source]

Bases: object

Partially directed graph after collider orientation.

Parameters:
class PCFGParseNode(label, span, log_prob, rule_index, rule_type, children=(), value=None)[source]

Bases: object

Node in a Viterbi parse tree.

Parameters:
leaves()[source]

Return terminal observations under this node.

Return type:

list[Any]

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

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])

  • noise (float)

  • residual_weight (float)

  • n_collocation (int)

  • m_steps (int)

  • lr (float)

  • seed (int)

  • name (str | None)

  • device (Any)

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

PINNRegressionEstimator

dist_to_encoder()[source]

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

Return type:

NeuralGaussianEncoder

to_dict()[source]

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

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Rebuild a PINNRegression from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

PINNRegression

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

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])

  • noise (float)

  • residual_weight (float)

  • n_collocation (int)

  • m_steps (int)

  • lr (float)

  • seed (int)

  • name (str | None)

  • device (Any)

accumulator_factory()[source]

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

Return type:

NeuralGaussianAccumulatorFactory

estimate(nobs, suff_stat)[source]

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

Parameters:
Return type:

PINNRegression

class QATWrapper(base, *, bits=4, clip_percentile=None, enabled=True)[source]

Bases: Module

Wrap an nn.Linear so its weight is straight-through fake-quantized on every forward call: the module computes F.linear(x, fake_quantize(weight), bias) instead of F.linear(x, weight, bias). Bias stays fp32, matching PTQ’s scheme (mixle.task.quantize.quantize_mlp()) where only weights are quantized.

Drop-in composition: QATWrapper(linear) has the same forward(x) -> Tensor contract as the Linear it wraps, so it slots into any module tree (see apply_qat()) without the surrounding model or training loop changing at all.

Parameters:

  • base (Any)

  • bits (int)

  • clip_percentile (float | None)

  • enabled (bool)

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Any)

Return type:

Any

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

apply_qat(model, *, bits=4, clip_percentile=None)[source]

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)

Return type:

Any

fake_quantize(x, *, bits=4, clip_percentile=None)[source]

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:
Return type:

Any

fake_quantize_int4(x, *, clip_percentile=None)[source]

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)[source]

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:
Return type:

Any

class PoissonRegressionNeuralNetwork(module, engine=None, precision=None)[source]

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()[source]

Return trainable parameters of the wrapped log-rate module.

Return type:

Iterable[Any]

log_rate_tensor(x)[source]

Return predicted log rates as a Torch tensor.

Parameters:

x (Any)

Return type:

Any

log_likelihood(x, y)[source]

Return the summed Poisson count log likelihood.

Parameters:
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)[source]

Maximize the Poisson count log likelihood.

Parameters:
Return type:

Any

predict_rate_tensor(x)[source]

Return predicted Poisson rates as a Torch tensor.

Parameters:

x (Any)

Return type:

Any

predict_rate(x)[source]

Return predicted Poisson rates as a NumPy array.

Parameters:

x (Any)

Return type:

ndarray

predict(x)[source]

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)[source]

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)

density(x)[source]

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

Parameters:

x (tuple[Any, Any])

Return type:

float

log_density(x)[source]

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

Parameters:

x (tuple[Any, Any])

Return type:

float

seq_log_density(x)[source]

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

Parameters:

x (tuple[ndarray, ndarray])

Return type:

ndarray

sample_y(x, rng)[source]

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

Parameters:
Return type:

ndarray

sampler(seed=None)[source]

Return a conditional sampler for drawing targets given features.

Parameters:

seed (int | None)

Return type:

RandomForestConditionalSampler

estimator(pseudo_count=None)[source]

Return a fresh estimator with the same task, name, and keyed-accumulation settings.

Parameters:

pseudo_count (float | None)

Return type:

RandomForestEstimator

dist_to_encoder()[source]

Return the encoder for feature/target observation pairs.

Return type:

RandomForestEncoder

class RandomForestEstimator(task='auto', 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)[source]

Bases: ParameterEstimator

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

task is ‘classification’, ‘regression’, or ‘auto’ (inferred from the dtype of y). 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)

  • 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)

accumulator_factory()[source]

Return an accumulator factory for weighted feature/target buffers.

Return type:

RandomForestAccumulatorFactory

estimate(nobs, suff_stat)[source]

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

Parameters:
Return type:

RandomForestConditional

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

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)[source]

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

Parameters:
Return type:

StochasticBlockGraphModel

log_likelihood(adjacency)[source]

Return the Bernoulli SBM log likelihood.

Parameters:

adjacency (Any)

Return type:

float

sample(seed=None)[source]

Draw one graph from the block model.

Parameters:

seed (int | None)

Return type:

ndarray

bic(adjacency)[source]

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)[source]

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)[source]

Create a randomly initialized model.

Parameters:
Return type:

TransEKnowledgeGraphModel

distance_triples(triples)[source]

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

Parameters:

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

Return type:

ndarray

score_triples(triples)[source]

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)[source]

Return the pairwise TransE ranking loss.

Parameters:
Return type:

float

negative_sample(triples, seed=None, corrupt='tail')[source]

Corrupt heads or tails to produce negative triples.

Parameters:
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)[source]

Fit embeddings with simple stochastic subgradient descent.

Parameters:
Return type:

KnowledgeGraphFitResult

normalize_entity_embeddings(max_norm=1.0)[source]

Project entity embeddings into an L2 ball.

Parameters:

max_norm (float)

Return type:

None

class TruncatedDirichletProcessMixtureFitResult(model, responsibilities, history)[source]

Bases: object

Fitted truncated DPM plus variational responsibilities and history.

Parameters:

  • model (TruncatedDirichletProcessMixtureModel)

  • responsibilities (ndarray)

  • history (list[float])

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

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)[source]

Return component log densities for one observation.

Parameters:

x (Any)

Return type:

ndarray

log_density(x)[source]

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

Parameters:

x (Any)

Return type:

float

density(x)[source]

Return the finite-truncation mixture density for one observation.

Parameters:

x (Any)

Return type:

float

responsibilities(data, expected=True)[source]

Return posterior component probabilities for observations.

Parameters:
Return type:

ndarray

effective_components(threshold=0.01)[source]

Count components with posterior mean stick weight above threshold.

Parameters:

threshold (float)

Return type:

int

sample(size=None, seed=None)[source]

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)[source]

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

Parameters:
Return type:

PartiallyObservableMarkovDecisionProcessFitResult

discrete_conditional_mutual_information(data, x, y, given=())[source]

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

Parameters:
Return type:

float

expected_log_stick_weights(gamma)[source]

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)[source]

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

Parameters:
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)[source]

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

Parameters:
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)[source]

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

The component M-steps are delegated to ordinary mixle.stats estimators. This keeps component likelihood math and sufficient statistics in their distribution modules.

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)[source]

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

Parameters:
Return type:

GrammarLearningResult

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

Fisher-z Gaussian conditional independence test.

Parameters:
Return type:

ConditionalIndependenceResult

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

Return partial correlation rho_xy.given for continuous data.

Parameters:
Return type:

float

grammar_rule_table(model)[source]

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)[source]

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

Parameters:
Return type:

HardEMResult

leapfrog_rollout(net, q0, p0, dt, n_steps)[source]

Symplectic 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. 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:
Return type:

tuple[Any, Any]

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

Learn a PC-style undirected skeleton from conditional independences.

Parameters:
Return type:

CausalSkeleton

build_autoregressive_categorical(dim, n_categories, *, hidden=64)[source]

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:
Return type:

Any

build_conditional_autoregressive_categorical(x_dim, y_dim, n_categories, *, hidden=64)[source]

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:
Return type:

Any

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

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:
Return type:

Any

build_convex_energy_net(dim, *, hidden=64, layers=3)[source]

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:
Return type:

Any

build_coupling_flow(dim, *, hidden=32, layers=4)[source]

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:
Return type:

Any

build_energy_net(dim, *, hidden=64, layers=3)[source]

An MLP energy E(x): R^dim -> R (plus a learned scalar log_norm) – ready to wrap in an EnergyModel.

Lower energy = higher (unnormalized) density. log_norm is the NCE-learned normalizer, so the paired EnergyModel scores -E(x) + log_norm. Swap in any module exposing energy(x) -> (n,), a log_norm parameter and a dim attribute.

Parameters:
Return type:

Any

build_maf(dim, *, hidden=64, blocks=3)[source]

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:
Return type:

Any

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

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:
Return type:

Any

build_product_energy_net(experts)[source]

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)[source]

A contrastive (InfoNCE / CLIP-style) conditional p(y | x) between two embedding spaces – ready to wrap.

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”.

log_density(x, y) returns, per row, the (negative) SYMMETRIC INFONCE loss for a batch of n paired embeddings: every row’s projected pair is scored against every OTHER row in the batch as a negative, in both directions (x -> y and y -> x), log-softmax-normalized over the batch dimension, then averaged. That is exactly what the shared NeuralConditionalDensity M-step already does with log_density – weight it and sum it – so no separate loss path is needed: the M-step’s responsibility-weighted-NLL gradient ascent on log_density is InfoNCE training, “for free” from the adapter’s existing contract. As with build_vae()’s ELBO, this is an honest score against itself (or another leaf scored the same batch-relative way) rather than a calibrated log p(y | x) – there is no way to integrate a softmax-over-the-current-batch score to 1 over all y. A batch of a single row has no negatives to contrast against, so log_density returns 0 for it rather than raising.

sample_given is not defined – a contrastive leaf is discriminative (it scores/ranks pairs); it has no generative p(y | x) to draw from. Retrieve a matching y by comparing module.embed_x(x) against module.embed_y(candidates) (cosine similarity in the shared space) instead.

Parameters:
Return type:

Any

build_vae(dim, *, latent=2, hidden=32)[source]

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.

log_density(x) returns the ELBO, a lower bound on log p(x), not the exact value. Compare VAE leaves with other bounded leaves whenever possible. Mixing a VAE with an exact-density leaf, such as a Gaussian or flow, compares a bound against an exact value and can under-weight the VAE.

log_density is deterministic: it evaluates the ELBO at the encoder mean z = mu(x) (no randn resample), so repeated scoring of the same x is bit-identical and an EM log-likelihood stays monotone. Training still uses the reparameterized sample (training=True) for an unbiased gradient.

Parameters:
Return type:

Any

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

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:
Return type:

Any

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

Create a simple fully connected Torch MLP.

Parameters:
Return type:

Any

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

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:
Return type:

Any

mean_stick_weights(gamma)[source]

Return E_q[pi_k] under independent Beta stick posteriors.

Parameters:

gamma (Any)

Return type:

ndarray

orient_v_structures(skeleton)[source]

Orient unshielded colliders i -> k <- j using separating sets.

Parameters:

skeleton (CausalSkeleton)

Return type:

PartiallyDirectedGraph

pcfg_log_likelihood(model, data)[source]

Return total PCFG log likelihood on raw sequences.

Parameters:
Return type:

float

sample_crp_assignments(num_obs, alpha, seed=None)[source]

Sample Chinese-restaurant-process assignments and table counts.

Parameters:
Return type:

tuple[ndarray, ndarray]

stick_breaking_weights(stick_fractions, residual=True)[source]

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)[source]

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

Parameters:

  • model (HeterogeneousPCFGDistribution)

  • sequence (Sequence[Any])

Return type:

PCFGParseNode

Submodules