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. The VAE reports a reproducible multi-sample estimate of its ELBO and exposes the estimate’s Monte Carlo standard error; it is not presented as an exact density. Any other density (a normalized energy model, …) plugs in the same way.

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

Bases: GradLeaf

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

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

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• precision (str)

• name (str | None)

log_density(x)

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

Parameters:

x (Any)

Return type:

float

seq_log_density(x)

Return per-row log densities for encoded observations.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

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

Parameters:

seed (int | None)

Return type:

NeuralDensitySampler

estimator(pseudo_count=None)

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

Parameters:

pseudo_count (float | None)

Return type:

NeuralDensityEstimator

dist_to_encoder()

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

Return type:

NeuralDensityEncoder

to_dict()

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

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a NeuralDensity from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralDensity

class NeuralDensitySampler(dist, seed=None)

Bases: DistributionSampler

Sampler for wrapped neural density modules exposing sample(n).

Parameters:

• dist (NeuralDensity)

• seed (int | None)

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

Draw observations from the wrapped module’s sampler.

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

class NeuralDensityEncoder

Bases: DataSequenceEncoder

Encode observations for vectorized neural-density scoring and fitting.

seq_encode(data)

Convert observations to a two-dimensional float array.

Parameters:

data (list)

Return type:

ndarray

class NeuralDensityAccumulator

Bases: SequenceEncodableStatisticAccumulator

Buffers the (responsibility-weighted) data for the M-step – the weights are the E-step’s soft counts.

update(x, weight, estimate)

Add one weighted observation to the accumulator.

Parameters:

• x (Any)

• weight (float)

• estimate (Any)

Return type:

None

seq_update(enc, weights, estimate)

Add an encoded batch and responsibility weights to the accumulator.

Parameters:

• enc (Any)

• weights (ndarray)

• estimate (Any)

Return type:

None

initialize(x, weight, rng)

Initialize from one observation using the ordinary update path.

Parameters:

• x (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 neural-density accumulator.

Parameters:

other (Any)

Return type:

NeuralDensityAccumulator

value()

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

Return type:

tuple

from_value(v)

Restore accumulator buffers from a value tuple.

Parameters:

v (tuple)

Return type:

NeuralDensityAccumulator

acc_to_encoder()

Return the encoder expected by this accumulator.

Return type:

NeuralDensityEncoder

class NeuralDensityAccumulatorFactory

Bases: StatisticAccumulatorFactory

Factory for neural-density accumulators.

make()

Create a fresh accumulator.

Return type:

NeuralDensityAccumulator

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

Bases: GradEstimator

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

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• precision (str)

• name (str | None)

accumulator_factory()

Return an accumulator factory for weighted neural-density batches.

Return type:

DataBufferAccumulatorFactory

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_vae(dim, *, latent=2, hidden=32, eval_samples=16)

Build a variational autoencoder over R^dim.

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

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

Parameters:

• dim (int)

• latent (int)

• hidden (int)

• eval_samples (int)

Return type:

Any

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

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

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

Parameters:

• dim (int)

• hidden (int)

• blocks (int)

Return type:

Any

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

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

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

Parameters:

• dim (int)

• n_categories (int)

• hidden (int)

Return type:

Any