mixle.models.neural_density module¶
NeuralDensity – the adapter that turns ANY torch density module into a composable mixle distribution.
The point is not a specific architecture; it is the wrapper. NeuralGaussian already adapts a conditional net
(p(y | x)); this is its unconditional sibling: give it any torch module that exposes log_density(x) -> (n,)
(and, to draw samples, sample(n) -> (n, d)) and you get a full five-piece mixle Distribution – so a
flexible neural density drops into a MixtureDistribution component, an HMM emission, or a
CompositeDistribution field, and is fit jointly with classical families by EM. Its M-step is a
responsibility-weighted maximum-likelihood gradient ascent on the module, warm-started across EM iterations.
That is the thing no NN library offers: “a mixture of a normalizing flow and a Gamma”, “an HMM whose emissions
are flows”. Ready instances ship: build_coupling_flow() (a RealNVP-style flow, exact), build_maf() (a masked
autoregressive flow, exact, richer autoregressive dependence), build_vae() (a variational autoencoder, a
latent-variable density whose log_density is the ELBO lower bound) and build_autoregressive_categorical()
(an exact autoregressive density over discrete vectors) – structurally different families, continuous and
discrete, behind one adapter. Any other density (a normalized energy model, …) plugs in the same way.
- class NeuralDensity(module, *, m_steps=60, lr=5e-3, device='cpu', name=None)[source]
Bases:
SequenceEncodableProbabilityDistributionWrap a torch density
module(module.log_density(x) -> (n,)) as a composable mixle distribution.- log_density(x)[source]
Return the log-density or log-mass at a single observation.
- seq_log_density(x)[source]
Return vectorized log-density values for sequence-encoded observations.
- sampler(seed=None)[source]
Return a sampler for drawing observations from this distribution.
- Parameters:
seed (int | None)
- Return type:
NeuralDensitySampler
- estimator(pseudo_count=None)[source]
Return an estimator for fitting this distribution from data.
- Parameters:
pseudo_count (float | None)
- Return type:
NeuralDensityEstimator
- dist_to_encoder()[source]
Return the data encoder used by this distribution for vectorized methods.
- Return type:
NeuralDensityEncoder
- to_dict()[source]
Return a safe JSON-compatible representation of this distribution.
- class NeuralDensitySampler(dist, seed=None)[source]
Bases:
DistributionSampler- Parameters:
dist (NeuralDensity)
seed (int | None)
- sample(size=None, *, batched=True)[source]
Draw observations.
Combinator samplers (mixture/sequence/…) accept
batched. Withbatched=True(the default) each child stream is drawn in one vectorized call instead of a per-draw Python loop – far faster. Because every child sampler owns an independentRandomState, batching consumes each stream in the same order as the loop, so the draws are identical to the legacy path.batched=Falseforces that legacy per-draw loop as a guaranteed- stable reference. Leaf samplers are already vectorized and ignore the flag.
- class NeuralDensityEncoder[source]
Bases:
DataSequenceEncoder
- class NeuralDensityAccumulator[source]
Bases:
SequenceEncodableStatisticAccumulatorBuffers the (responsibility-weighted) data for the M-step – the weights are the E-step’s soft counts.
- update(x, weight, estimate)[source]
- seq_update(enc, weights, estimate)[source]
- seq_initialize(enc, weights, rng)[source]
- acc_to_encoder()[source]
- Return type:
NeuralDensityEncoder
- class NeuralDensityAccumulatorFactory[source]
Bases:
StatisticAccumulatorFactory- make()[source]
- Return type:
NeuralDensityAccumulator
- class NeuralDensityEstimator(module, *, m_steps=60, lr=5e-3, device='cpu', name=None)[source]
Bases:
ParameterEstimatorM-step: responsibility-weighted MLE –
max sum_i w_i log p(x_i)by gradient ascent on the module (warm).- accumulator_factory()[source]
- Return type:
NeuralDensityAccumulatorFactory
- build_coupling_flow(dim, *, hidden=32, layers=4)[source]
A RealNVP coupling flow over
R^dimwith an exactlog_density(x)andsample(n)– ready to wrap.Alternating affine-coupling layers map data to a standard-normal base;
log_densityis the base log-prob plus the log-determinant of the (triangular) Jacobian. A minimal, correct instance of the density module aNeuralDensityadapts – swap in any other module with the same two methods.
- build_vae(dim, *, latent=2, hidden=32)[source]
A variational autoencoder over
R^dim– a latent-variable densityp(x) = int p(x | z) p(z) dz.An amortized encoder
q(z | x)and a decoderp(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 sameNeuralDensityadapter, because it exposes the same two methods.Caveat, stated plainly:
log_density(x)returns the ELBO, a lower bound onlog 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_densityis deterministic: it evaluates the ELBO at the encoder meanz = mu(x)(norandnresample), so repeated scoring of the samexis bit-identical and an EM log-likelihood stays monotone. Training still uses the reparameterized sample (training=True) for an unbiased gradient.
- build_maf(dim, *, hidden=64, blocks=3)[source]
A masked autoregressive flow over
R^dim– an exact density that factorizesp(x)by the chain rule, eachp(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_densityis 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 forNeuralDensity; the adapter is unchanged.
- build_autoregressive_categorical(dim, n_categories, *, hidden=64)[source]
An autoregressive neural density over discrete vectors
x in {0..C-1}^dim– exact, normalizedp(x).The continuous flows/VAE above model
R^d; heterogeneous data is also categorical. This factorizesp(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_densitysums the picked log-softmax logits;samplefills the vector one coordinate at a time. Another ready module forNeuralDensity; the adapter is unchanged.