mixle.models.mixture_density module

NeuralConditionalDensity – the adapter that turns ANY torch conditional density into a mixle leaf.

This is the conditional sibling of NeuralDensity. Where that one wraps a module exposing log_density(x) -> (n,) (an unconditional p(x)), this wraps a module exposing log_density(x, y) -> (n,) (and sample_given(x) -> (n, d)) and gives you a full five-piece mixle Distribution over the pair (x, y) – so a flexible conditional density drops into a mixture of experts, a composite field, or an HMM emission and is fit jointly with classical families by the same responsibility-weighted-NLL EM M-step (warm-started across iterations, i.e. generalized EM).

Why it matters: NeuralGaussian fixes the conditional law to a single Gaussian, p(y | x) = N(y; f(x), sigma^2 I) – one mean per x, unimodal and homoscedastic. Many real conditionals are neither: an inverse problem has several valid y for one x; measurement noise grows with x. build_mdn() is the ready instance – a mixture density network, p(y | x) = sum_k pi_k(x) N(y; mu_k(x), sigma_k(x)^2) – whose entire mixture (weights, means, variances) is a function of x, so it is multimodal and heteroscedastic. Any other conditional density (a conditional flow, an autoregressive head) plugs in the same way: give it log_density(x, y) and sample_given(x).

build_projection_leaf() is a different kind of ready instance – a contrastive p(y | x) (an InfoNCE projection between two, typically frozen, embedding spaces) whose log_density is not a calibrated density at all but still trains and composes through the exact same adapter: the stage-1 “frozen encoder -> projection -> frozen encoder” pattern, generalized to a family with no domain nouns.

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

Bases: SequenceEncodableProbabilityDistribution

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

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

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• name (str | None)

log_density(xy)

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

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)

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

Parameters:

enc (Any)

Return type:

ndarray

sampler(seed=None)

Return a conditional sampler for drawing y given x.

Parameters:

seed (int | None)

Return type:

NeuralConditionalDensitySampler

estimator(pseudo_count=None)

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

Parameters:

pseudo_count (float | None)

Return type:

NeuralConditionalDensityEstimator

dist_to_encoder()

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

Return type:

NeuralConditionalDensityEncoder

to_dict()

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

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a NeuralConditionalDensity from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralConditionalDensity

class NeuralConditionalDensitySampler(dist, seed=None)

Bases: DistributionSampler

Conditional sampler for modules exposing sample_given(x).

Parameters:

• dist (NeuralConditionalDensity)

• seed (int | None)

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

Raise because the leaf defines p(y | x) and has no marginal p(x).

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

sample_given(x)

Draw one response from p(y | x) using the wrapped module.

Parameters:

x (Any)

Return type:

ndarray

sample_given_batch(x_batch)

One draw of y ~ p(y | x) for every row of x_batch (shape (n, x_dim)), in one batched forward pass – statistically identical to calling sample_given() once per row (same model, same per-draw sampling procedure), just without paying framework/dispatch overhead per row. That per-call overhead dominates a Python loop of hundreds of individual sample_given calls, which is exactly the shape both a particle-walk step (many different x’s, one draw each) and a per-point coverage check (repeat one x, many draws) reduce to – both call sites use this to speed up the same check/walk rather than shrink it. Repeat a row of x_batch to draw more than once from the same x.

Parameters:

x_batch (Any)

Return type:

ndarray

class NeuralConditionalDensityEncoder

Bases: DataSequenceEncoder

Encode (x, y) pairs for vectorized conditional-density scoring and fitting.

seq_encode(data)

Convert a list of (x, y) pairs into batched feature and target arrays.

Parameters:

data (list)

Return type:

tuple[ndarray, ndarray]

class NeuralConditionalDensityAccumulator

Bases: SequenceEncodableStatisticAccumulator

Buffers responsibility-weighted (x, y) pairs for the M-step (the weights are the E-step soft counts).

update(xy, weight, estimate)

Add one weighted observation pair to the accumulator.

Parameters:

• xy (Any)

• weight (float)

• estimate (Any)

Return type:

None

seq_update(enc, weights, estimate)

Add a batch of encoded observation pairs and responsibility weights.

Parameters:

• enc (Any)

• weights (ndarray)

• estimate (Any)

Return type:

None

initialize(xy, weight, rng)

Initialize from one observation using the ordinary update path.

Parameters:

• xy (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

Initialize from an encoded batch using the ordinary batch update path.

Parameters:

• enc (Any)

• weights (ndarray)

• rng (Any)

Return type:

None

combine(other)

Merge the value tuple from another conditional-density accumulator.

Parameters:

other (Any)

Return type:

NeuralConditionalDensityAccumulator

value()

Return contiguous (x, y, weights) arrays for the M-step.

Return type:

tuple

from_value(value)

Restore accumulator buffers from a value tuple.

Parameters:

value (tuple)

Return type:

NeuralConditionalDensityAccumulator

acc_to_encoder()

Return the encoder expected by this accumulator.

Return type:

NeuralConditionalDensityEncoder

class NeuralConditionalDensityAccumulatorFactory

Bases: StatisticAccumulatorFactory

Factory for conditional-density accumulators.

make()

Create a fresh accumulator.

Return type:

NeuralConditionalDensityAccumulator

class NeuralConditionalDensityEstimator(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(y_i | x_i) by gradient ascent (warm-started).

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• name (str | None)

accumulator_factory()

Return an accumulator factory for weighted conditional-density batches.

Return type:

NeuralConditionalDensityAccumulatorFactory

estimate(nobs, suff_stat)

Run the weighted conditional log-likelihood M-step and return the updated leaf.

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

NeuralConditionalDensity

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_conditional_flow(x_dim, y_dim, *, hidden=32, layers=4)

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

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

Parameters:

• x_dim (int)

• y_dim (int)

• hidden (int)

• layers (int)

Return type:

Any

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

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:

• d_x (int)

• d_y (int)

• encoder_x (Any)

• encoder_y (Any)

• proj_dim (int | None)

• hidden (int)

• freeze_encoders (bool)

• temperature (float)

Return type:

Any

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

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

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

Parameters:

• x_dim (int)

• y_dim (int)

• n_categories (int)

• hidden (int)

Return type:

Any