mixle.models.softmax_leaf module

A neural classifier as a mixle conditional-density leaf: p(y | x) = softmax(module(x)).

The discriminative sibling of NeuralGaussian. NeuralCategorical(module) wraps a Torch module that emits k logits as a mixle distribution over observations (x, y) with y an integer class index. It implements the full SequenceEncodableProbabilityDistribution contract, so it drops into MixtureDistribution / CompositeDistribution / HMM emissions like any leaf – and its EM M-step is a responsibility-weighted cross-entropy gradient step on the module (warm-started across EM iterations => generalized EM). The model’s seq_log_density IS -cross_entropy(module(x), y): the objective is the leaf’s log-density, never a user-supplied loss closure.

This is the leaf that the declarative Categorical(logits=Net(...)) PPL slot lowers to, and the component that makes a Mix([Categorical(logits=Net(...)), ...]) a mixture of neural classifiers fit by ordinary EM.

Requires torch. The leaf is conditional: predict(x) / sampler().sample_given(x) work; sample() raises (there is no p(x)) – the same honest limitation NeuralGaussian and RandomForestConditional carry.

class NeuralCategorical(module, m_steps=40, lr=0.01, name=None, batch_size=None, device='cpu')[source]

Bases: SequenceEncodableProbabilityDistribution

p(y | x) = softmax(module(x)) as a mixle leaf. Observation is the pair (x, y), y an int class.

batch_size (None = full batch) makes the M-step minibatch SGD over m_steps passes – needed to train a real conv net on a large image set; device (e.g. "mps"/"cuda") runs it on the GPU.

Parameters:

  • module (Any)

  • m_steps (int)

  • lr (float)

  • name (str | None)

  • batch_size (int | None)

  • device (str)

log_density(xy)[source]

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

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)[source]

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

Parameters:

enc (Any)

Return type:

ndarray

predict(x)[source]
Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

NeuralCategoricalSampler

estimator(pseudo_count=None)[source]

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

NeuralCategoricalEstimator

dist_to_encoder()[source]

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

Return type:

NeuralCategoricalEncoder

to_dict()[source]

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralCategorical

class NeuralCategoricalSampler(dist, seed=None)[source]

Bases: DistributionSampler

Parameters:

  • dist (NeuralCategorical)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

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

Any

sample_given(x)[source]
Parameters:

x (Any)

Return type:

int

class NeuralCategoricalEncoder[source]

Bases: DataSequenceEncoder

seq_encode(data)[source]

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

data (list)

Return type:

tuple[ndarray, ndarray]

class NeuralCategoricalAccumulator[source]

Bases: SequenceEncodableStatisticAccumulator

update(xy, weight, estimate)[source]
Parameters:
Return type:

None

seq_update(enc, weights, estimate)[source]
Parameters:
Return type:

None

initialize(xy, weight, rng)[source]
Parameters:
Return type:

None

seq_initialize(enc, weights, rng)[source]
Parameters:
Return type:

None

combine(other)[source]
Parameters:

other (Any)

Return type:

NeuralCategoricalAccumulator

value()[source]
Return type:

tuple

from_value(value)[source]
Parameters:

value (tuple)

Return type:

NeuralCategoricalAccumulator

acc_to_encoder()[source]
Return type:

NeuralCategoricalEncoder

class NeuralCategoricalAccumulatorFactory[source]

Bases: StatisticAccumulatorFactory

make()[source]
Return type:

NeuralCategoricalAccumulator

class NeuralCategoricalEstimator(module, m_steps=40, lr=0.01, name=None, batch_size=None, device='cpu', ewc=None)[source]

Bases: ParameterEstimator

EM estimator for a NeuralCategorical: the M-step is m_steps of responsibility-weighted cross-entropy gradient on the module (the module is warm-started across EM iterations => generalized EM).

The weighted CE is normalized by the responsibility mass sum(w) so the M-step is scale-invariant to the responsibility magnitude (the easy bug: an unnormalized weighted loss makes the step size track cluster size).

Parameters:

  • module (Any)

  • m_steps (int)

  • lr (float)

  • name (str | None)

  • batch_size (int | None)

  • device (str)

  • ewc (Any)

accumulator_factory()[source]
Return type:

NeuralCategoricalAccumulatorFactory

estimate(nobs, suff_stat)[source]
Parameters:
Return type:

NeuralCategorical

SoftmaxNeuralLeaf

alias of NeuralCategorical

SoftmaxNeuralLeafEstimator

alias of NeuralCategoricalEstimator