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 trained LM to path via torch.save (hyperparameters + weights).

Parameters:

path (str)

Return type:

None

classmethod load(path)

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)

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

Parameters:

• token_ids (Any)

• epochs (int)

• batch_size (int)

• lr (float)

• distributed (bool)

• precision (str)

• shuffle (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 CategoricalEmbedding(num_categories, dim, *, name=None)

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

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

Return type:

Any

class CausalSkeleton(edges, separating_sets, variable_names)

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

edges: set[tuple[int, int]]

separating_sets: dict[tuple[int, int], frozenset[int]]

variable_names: list[Any]

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)

Bases: object

Result from a conditional independence calculation.

Parameters:

• measure (float)

• statistic (float)

• p_value (float | None)

• independent (bool)

measure: float

statistic: float

p_value: float | None

independent: bool

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

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_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 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')

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)

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

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)

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

Parameters:

enc (Any)

Return type:

ndarray

predict(x)

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

NeuralCategoricalSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

NeuralCategoricalEstimator

dist_to_encoder()

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

Return type:

NeuralCategoricalEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

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 the log-density or log-mass at a single observation.

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)

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

Parameters:

enc (Any)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

NeuralConditionalDensitySampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

NeuralConditionalDensityEstimator

dist_to_encoder()

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

Return type:

NeuralConditionalDensityEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralConditionalDensity

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

Bases: SequenceEncodableProbabilityDistribution

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

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• name (str | 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:

NeuralDensitySampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

NeuralDensityEstimator

dist_to_encoder()

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

Return type:

NeuralDensityEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralDensity

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

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

Return type:

NeuralDensityAccumulatorFactory

estimate(nobs, suff_stat)

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

NeuralDensity

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 the log-density or log-mass at a single observation.

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)

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

Parameters:

enc (Any)

Return type:

ndarray

classmethod compute_capabilities()

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 sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

NeuralGaussianSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

NeuralGaussianEstimator

dist_to_encoder()

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

Return type:

NeuralGaussianEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralGaussian

class VAE(dim, *, latent=2, hidden=32, 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 the ELBO – a lower bound on log p(x), evaluated deterministically at the encoder mean (so an EM log-likelihood stays monotone). Honest on its own, in a mixture of VAEs, or against another bounded leaf; mixing it with an exact-density leaf (a Gaussian, a flow) compares a bound against an exact value and under-weights the VAE. See build_vae() for the full statement.

Parameters:

• dim (int)

• latent (int)

• hidden (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 vectorized log-density values for sequence-encoded observations.

Parameters:

enc (Any)

Return type:

ndarray

log_density(xcr)

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

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 a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

DPOModelSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

DPOModelEstimator

dist_to_encoder()

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

Return type:

DPOEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

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

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:

EnergyModelSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

EnergyModelEstimator

dist_to_encoder()

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

Return type:

EnergyModelEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

EnergyModel

class StreamingTransformer(module, device='cpu')

Bases: SequenceEncodableProbabilityDistribution

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

Parameters:

• module (Any)

• device (str)

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

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

embedding optionally ties a shared CategoricalEmbedding across leaves.

Parameters:

• vocab (int)

• d_model (int)

• n_layer (int)

• n_head (int)

• block (int)

• embedding (Any)

• device (str)

Return type:

StreamingTransformer

log_density(xy)

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

Parameters:

xy (Any)

Return type:

float

predict(x)

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

StreamingTransformerSampler

seq_log_density(enc)

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

Parameters:

enc (Any)

Return type:

ndarray

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

StreamingTransformerEstimator

dist_to_encoder()

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

Return type:

StreamingTokenEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

StreamingTransformer

NeuralLeaf

alias of NeuralGaussian

SoftmaxNeuralLeaf

alias of NeuralCategorical

StreamingTransformerLeaf

alias of StreamingTransformer

class TrainingSearchResult(recipe, loss, history=None)

Bases: object

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

Parameters:

• recipe (dict[str, Any])

• loss (float)

• history (Any)

recipe: dict[str, Any]

loss: float

history: Any = 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])

d_model_choices: Sequence[int] = (64, 128, 256, 512)

n_layer_range: tuple[int, int] = (2, 12)

log10_lr_range: tuple[float, float] = (-4.0, -2.0)

batch_choices: Sequence[int] = (16, 32, 64, 128)

dims()

Return type:

int

bounds()

Return type:

list[tuple[float, float]]

decode(point)

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)

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

Parameters:

• vocab (int)

• d_model (int)

• n_layer (int)

• n_head (int)

• block (int)

• embedding (Any)

Return type:

Any

ewc(anchor, fisher, lam=1.0)

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

Parameters:

• anchor (list)

• fisher (list)

• lam (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) (the standard learning-curve form) and evaluates it at at – so a cheap partial run can estimate the full-budget loss for early stopping. Falls back to the last observed loss if the fit fails.

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 of a classification leaf’s module on (x, y): mean of (d log p(y|x)/dtheta)^2.

Parameters:

• leaf (Any)

• x (Any)

• y (Any)

• samples (int)

• device (str)

• seed (int)

Return type:

list

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

A ready training callback (recipe, budget) -> held-out nats/token that trains a real mixle LM.

budget in (0, 1] scales the number of epochs (the cheap-fidelity axis is training length); a real pretraining loop would scale steps or the token subset the same way.

Parameters:

• token_ids (Sequence[int])

• val_ids (Sequence[int])

• vocab (int)

• block (int)

• max_epochs (int)

• device (str)

Return type:

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

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:

list

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

Multi-fidelity BO of the training recipe; train(recipe, budget) returns held-out loss (lower is better).

fidelities are the training-budget fractions the search may run at (cheap first). Returns the recipe with the best full-budget loss and the full BO history. Scale the objective by swapping train for a real launcher – the search machinery is identical.

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

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

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

Parameters:

• module (Any)

• token_source (Any)

• lr (float)

• device (str)

• report_every (int)

• log (Any)

Return type:

tuple

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

Bases: object

Belief trajectories, log likelihood, and predictive observation terms.

Parameters:

• beliefs (ndarray)

• log_likelihood (float)

• predictive_observation_probs (ndarray)

beliefs: ndarray

log_likelihood: float

predictive_observation_probs: ndarray

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

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)

label: Any

span: tuple[int, int]

log_prob: float

rule_index: int

rule_type: str

children: tuple[PCFGParseNode, ...] = ()

value: Any = None

leaves()

Return terminal observations under this node.

Return type:

list[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)

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)

Return the probability density or mass at a single observation.

Concrete default: exponentiate log_density (the abstract method subclasses must provide). Leaves with a cheaper closed form may override this.

Parameters:

x (tuple[Any, Any])

Return type:

float

log_density(x)

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

Parameters:

x (tuple[Any, Any])

Return type:

float

seq_log_density(x)

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

Parameters:

x (tuple[ndarray, ndarray])

Return type:

ndarray

sample_y(x, rng)

Parameters:

• x (Any)

• rng (RandomState)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

RandomForestConditionalSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

RandomForestEstimator

dist_to_encoder()

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

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)

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

Return type:

RandomForestAccumulatorFactory

estimate(nobs, suff_stat)

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

Corrupt heads or tails to produce negative triples.

Parameters:

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

• seed (int | None)

• corrupt (str)

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)

Bases: object

Fitted truncated DPM plus variational responsibilities and history.

Parameters:

• model (TruncatedDirichletProcessMixtureModel)

• responsibilities (ndarray)

• history (list[float])

model: TruncatedDirichletProcessMixtureModel

responsibilities: ndarray

history: list[float]

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

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

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 composes honestly in a mixture 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_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 composes honestly. 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 (so it composes honestly, unlike 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_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 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:

• 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 – so it composes honestly in a mixture with a Gaussian, a flow, or any exact 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_vae(dim, *, latent=2, hidden=32)

A variational autoencoder over R^dim – a latent-variable density p(x) = int p(x | z) p(z) dz.

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. It is a genuinely different family from the flow – structure through a low-dimensional latent, not an invertible map – yet it plugs into the same NeuralDensity adapter, because it exposes the same two methods.

Caveat, stated plainly: log_density(x) returns the ELBO, a lower bound on log p(x), not the exact value (the flow’s is exact). So a VAE leaf is honest on its own, in a mixture of VAEs, or against another bounded leaf – but mixing it with an exact-density leaf (a Gaussian, a flow) compares a bound against an exact value and will under-weight the VAE. Because ELBO <= log p(x), a VAE that beats an exact leaf on held-out data still wins by at least that margin; a VAE that loses may not actually be worse.

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:

• dim (int)

• latent (int)

• hidden (int)

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

mean_stick_weights(gamma)

Return E_q[pi_k] under independent Beta stick posteriors.

Parameters:

gamma (Any)

Return type:

ndarray

orient_v_structures(skeleton)

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

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.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.gaussian_process module
• mixle.models.grammar module
• mixle.models.knowledge_graph module
• mixle.models.language_model module
• mixle.models.mixture_density module
• mixle.models.neural module
• mixle.models.neural_density module
• mixle.models.neural_families module
• mixle.models.neural_leaf module
• mixle.models.partially_observable_markov_decision_process module
• mixle.models.random_forest module
• mixle.models.random_graph module
• mixle.models.softmax_leaf module
• mixle.models.sparse_gaussian_process module
• mixle.models.streaming_transformer_leaf module
• mixle.models.train_search module
• mixle.models.transformer module