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 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 NeuralGaussianSampler(dist, seed=None)

Bases: DistributionSampler

Parameters:

• dist (NeuralGaussian)

• seed (int | None)

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

Draw observations.

Combinator samplers (mixture/sequence/…) accept batched. With batched=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 independent RandomState, batching consumes each stream in the same order as the loop, so the draws are identical to the legacy path. batched=False forces that legacy per-draw loop as a guaranteed- stable reference. Leaf samplers are already vectorized and ignore the flag.

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

sample_given(x)

Parameters:

x (Any)

Return type:

ndarray

class NeuralGaussianEncoder

Bases: DataSequenceEncoder

seq_encode(data)

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

data (list)

Return type:

tuple[ndarray, ndarray]

class NeuralGaussianAccumulator

Bases: SequenceEncodableStatisticAccumulator

update(xy, weight, estimate)

Parameters:

• xy (Any)

• weight (float)

• estimate (Any)

Return type:

None

seq_update(enc, weights, estimate)

Parameters:

• enc (Any)

• weights (ndarray)

• estimate (Any)

Return type:

None

initialize(xy, weight, rng)

Parameters:

• xy (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

Parameters:

• enc (Any)

• weights (ndarray)

• rng (Any)

Return type:

None

combine(other)

Parameters:

other (Any)

Return type:

NeuralGaussianAccumulator

value()

Return type:

tuple

from_value(value)

Parameters:

value (tuple)

Return type:

NeuralGaussianAccumulator

acc_to_encoder()

Return type:

NeuralGaussianEncoder

class NeuralGaussianAccumulatorFactory

Bases: StatisticAccumulatorFactory

make()

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 type:

NeuralGaussianAccumulatorFactory

estimate(nobs, suff_stat)

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

NeuralGaussian

NeuralLeaf

alias of NeuralGaussian

NeuralLeafEstimator

alias of NeuralGaussianEstimator