mixle.models.neural_leaf module

A neural network as a mixle conditional-density leaf – the bridge that makes nets generative components.

NeuralGaussian(module) wraps a Torch module as a mixle distribution p(y | x) = N(y; module(x), noise^2 I) over observations (x, y). It implements the full SequenceEncodableProbabilityDistribution contract, so it drops into MixtureDistribution / CompositeDistribution / HMM emissions like any leaf – but its EM M-step is weighted-NLL gradient descent on the module (warm-started across EM iterations => generalized EM).

A MixtureDistribution of NeuralGaussian components is therefore a mixture of neural experts: the E-step computes responsibilities, the M-step trains each expert by responsibility-weighted regression. Combined with the em move in mixle.experimental.program, the same model fits with EM where conjugate and gradient where neural:

from mixle.stats import MixtureEstimator
experts = MixtureEstimator([NeuralGaussian(mlp_a).estimator(), NeuralGaussian(mlp_b).estimator()])
# ... run EM (estimate loop) -> each expert specializes, gated by the responsibilities.

Requires torch (the module). The leaf is conditional: sampler().sample_given(x) draws y; sample() raises (there is no p(x)).

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 log p(y | x) for one encoded observation pair (x, y).

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)

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

Parameters:

enc (Any)

Return type:

ndarray

classmethod compute_capabilities()

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

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 conditional sampler for drawing y given x.

Parameters:

seed (int | None)

Return type:

NeuralGaussianSampler

estimator(pseudo_count=None)

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

Parameters:

pseudo_count (float | None)

Return type:

NeuralGaussianEstimator

dist_to_encoder()

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

Return type:

NeuralGaussianEncoder

to_dict()

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

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a NeuralGaussian from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralGaussian

class NeuralGaussianSampler(dist, seed=None)

Bases: DistributionSampler

Conditional sampler for NeuralGaussian; draws responses given covariates.

Parameters:

• dist (NeuralGaussian)

• 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 Gaussian response from p(y | x).

Parameters:

x (Any)

Return type:

ndarray

class NeuralGaussianEncoder

Bases: DataSequenceEncoder

Encode (x, y) observation pairs for vectorized neural-Gaussian 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 NeuralGaussianAccumulator

Bases: SequenceEncodableStatisticAccumulator

Buffer weighted (x, y) batches for the neural-Gaussian M-step.

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 neural-Gaussian accumulator.

Parameters:

other (Any)

Return type:

NeuralGaussianAccumulator

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:

NeuralGaussianAccumulator

acc_to_encoder()

Return the encoder expected by this accumulator.

Return type:

NeuralGaussianEncoder

class NeuralGaussianAccumulatorFactory

Bases: StatisticAccumulatorFactory

Factory for neural-Gaussian accumulators.

make()

Create a fresh accumulator.

Return type:

NeuralGaussianAccumulator

class NeuralGaussianEstimator(module, noise=1.0, m_steps=40, lr=0.01, name=None, device=None)

Bases: ParameterEstimator

EM estimator for a NeuralGaussian: the M-step is m_steps of weighted-NLL gradient on the module.

The module is held (and warm-started) across EM iterations, so each M-step is a partial maximization (generalized EM). The accumulator buffers responsibility-weighted (x, y) observations.

Parameters:

• module (Any)

• noise (float)

• m_steps (int)

• lr (float)

• name (str | None)

• device (Any)

accumulator_factory()

Return an accumulator factory for weighted neural-regression batches.

Return type:

NeuralGaussianAccumulatorFactory

estimate(nobs, suff_stat)

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

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

NeuralGaussian

NeuralLeaf

alias of NeuralGaussian

NeuralLeafEstimator

alias of NeuralGaussianEstimator