mixle.models.grad_leaf module

GradLeaf – a torch module IS the model: the five-piece contract, manufactured.

The contract (Distribution / Sampler / Estimator / Accumulator / DataEncoder) earns its keep for closed-form families: additive sufficient statistics are what make EM exact and distributable. A GRADIENT leaf has no sufficient statistics – its “accumulator” can only buffer the responsibility-weighted data, its encoder is np.asarray, and its M-step is SGD – so per-family contract code is pure ceremony (mixle.models grew nine hand-written buffer accumulators saying so). This module writes that ceremony ONCE, generically. A neural family is now just a module:

fitted = optimize(x, module) # a bare nn.Module coerces – no wrapper at all leaf = GradLeaf(module) # or wrap explicitly to set knobs/hooks … mix = MixtureDistribution([leaf, gamma], w) # … and compose with classical families

The module owns forward and objective; mixle owns the loop. The contract’s requirements on the module are two methods: log_density(x) -> (n,) (scoring; also the default M-step objective) and, only if you draw samples, sample(n) -> (n, d). Control never leaves the caller:

• loss(module, x, w) -> scalar overrides the default responsibility-weighted NLL – custom objectives are a hook, not a subclass tree;

• optimizer(params) -> torch.optim.Optimizer picks the optimizer; it receives only TRAINABLE parameters, so freezing submodules with requires_grad_(False) just works (train a projection head against a frozen encoder; a FULLY frozen module is a fixed distribution and the M-step is a no-op);

• fitted.module is the raw torch module – nothing is trapped.

Serialization: the module round-trips as portable bytes (mixle.models._neural_serial); custom loss/optimizer hooks must be module-level functions to survive pickling, like any hook.

class DataBufferAccumulator(encoder, n_fields=1)

Bases: SequenceEncodableStatisticAccumulator

THE gradient-leaf “sufficient statistic”: the encoded, responsibility-weighted data itself, buffered for the M-step (the weights are the E-step’s soft counts). Generic over the encoding arity – a single array for unconditional leaves, a tuple like (x, y) for conditional ones – so every gradient family shares this one class instead of hand-writing its own buffer. Single observations route through the family’s own encoder, so per-row quirks live in exactly one place.

Parameters:

• encoder (Any)

• n_fields (int)

update(x, weight, estimate)

Accumulate one weighted observation under an optional current estimate.

Parameters:

• x (Any)

• weight (float)

• estimate (Any)

Return type:

None

seq_update(enc, weights, estimate)

Accumulate weighted sufficient statistics from sequence-encoded observations.

Parameters:

• enc (Any)

• weights (ndarray)

• estimate (Any)

Return type:

None

initialize(x, weight, rng)

Initialize sufficient statistics from one weighted observation.

Parameters:

• x (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

Initialize sufficient statistics from sequence-encoded observations.

Parameters:

• enc (Any)

• weights (ndarray)

• rng (Any)

Return type:

None

combine(other)

Merge serialized sufficient statistics into this accumulator.

Parameters:

other (Any)

Return type:

DataBufferAccumulator

value()

Return this accumulator’s serialized sufficient statistics.

Return type:

tuple

from_value(v)

Restore this accumulator from serialized sufficient statistics.

Parameters:

v (tuple)

Return type:

DataBufferAccumulator

acc_to_encoder()

Return a sequence encoder compatible with this accumulator.

Return type:

Any

class DataBufferAccumulatorFactory(encoder, n_fields=1)

Bases: StatisticAccumulatorFactory

Parameters:

• encoder (Any)

• n_fields (int)

make()

Create a fresh accumulator instance.

Return type:

DataBufferAccumulator

class GradEstimator(module, *, m_steps=60, lr=5e-3, device=None, batch_size=None, precision='fp32', name=None, loss=None, optimizer=None)

Bases: ParameterEstimator

M-step: responsibility-weighted MLE – max sum_i w_i log p(x_i) by gradient ascent on the module (warm-started across EM iterations). loss/optimizer are the caller’s hooks; the optimizer only ever sees trainable parameters, so frozen submodules stay frozen and a fully frozen module makes the M-step a no-op (a fixed distribution).

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (Any)

• batch_size (int | None)

• precision (str)

• name (str | None)

• loss (Any)

• optimizer (Any)

accumulator_factory()

Return the accumulator factory used to collect this estimator’s sufficient statistics.

Return type:

DataBufferAccumulatorFactory

estimate(nobs, suff_stat)

Estimate a distribution from accumulated sufficient statistics.

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

GradLeaf

class GradLeaf(module, *, m_steps=60, lr=5e-3, device=None, batch_size=None, precision='fp32', name=None, loss=None, optimizer=None)

Bases: SequenceEncodableProbabilityDistribution

Wrap a torch density module (module.log_density(x) -> (n,)) as a composable mixle distribution (see the module docstring). loss and optimizer are the M-step hooks.

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (Any)

• batch_size (int | None)

• precision (str)

• name (str | None)

• loss (Any)

• optimizer (Any)

log_density(x)

Return the log-density or log-mass at a single observation.

Parameters:

x (Any)

Return type:

float

seq_log_density(x)

Return vectorized log-density values for sequence-encoded observations.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

GradLeafSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

GradEstimator

dist_to_encoder()

Return the data encoder used by this distribution for vectorized methods.

Return type:

GradLeafEncoder