mixle.stats.compute.stacked module

Stacked mixture kernels built from distribution-owned backend math.

class StackedComponentParams(component_type, strategy, params)[source]

Bases: object

A generic route for homogeneous component scoring.

Parameters:
component_type: type[Any]
strategy: str
params: Any
class StackedEstimatorView(estimators)[source]

Bases: object

Minimal estimator view for component-stacked resident reductions.

Parameters:

estimators (Sequence[Any])

estimators: Sequence[Any]
class StackedMixtureResidentStats(component_counts, component_stats, engine, component_type)[source]

Bases: object

Engine-resident sufficient statistics for a homogeneous mixture.

Parameters:

  • component_counts (Any)

  • component_stats (Any)

  • engine (ComputeEngine)

  • component_type (type[Any])

component_counts: Any
component_stats: Any
engine: ComputeEngine
component_type: type[Any]
value()[source]

Return sufficient statistics in the legacy MixtureEstimator format.

Return type:

Any

local_value()[source]

Return this rank’s component-stat payload in the legacy local-slice format.

DTensor-backed model-parallel runs should run component M-steps on the local component shard and all-reduce only scalar mixture-weight totals. This method therefore prefers DTensor to_local() chunks over full_tensor() and returns (component_start, (counts, stats)). Non-sharded engines naturally return the full component range starting at zero.

Return type:

tuple[int, Any]

estimate(estimator)[source]

Estimate a distribution after converting to the legacy estimator protocol.

Parameters:

estimator (ParameterEstimator)

Return type:

Any

estimate_component_shard(estimator, total_count=None)[source]

Run the component-local part of a mixture M-step for this shard.

Component distributions are independent given posterior sufficient statistics. Mixture weights only need the scalar global total count (or fixed weights), so a model-parallel orchestrator can call this on each shard after an all-reduce of that scalar instead of gathering all component statistics to the driver.

Parameters:

  • estimator (ParameterEstimator)

  • total_count (float | None)

Return type:

StackedMixtureShardEstimate

class StackedMixtureShardEstimate(component_start, component_stop, components, weights, total_count)[source]

Bases: object

Component-local M-step result for a homogeneous mixture shard.

Parameters:
component_start: int
component_stop: int
components: tuple[Any, ...]
weights: ndarray
total_count: float
estimate_component_shard_value(estimator, component_start, value, total_count=None)[source]

Estimate the component-local part of a mixture M-step from an explicit shard value.

Parameters:
Return type:

StackedMixtureShardEstimate

tie_component_shard_values(estimator, shard_values)[source]

Apply mixle key tying to component-sharded mixture statistics.

MixtureAccumulator.key_merge / key_replace assume a full component vector is materialized on one worker. Component-sharded model-parallel runs instead hold ranges [component_start, component_stop). This helper applies the same key protocol to those local ranges by materializing only the component accumulators owned by each shard.

Parameters:
Return type:

tuple[tuple[int, tuple[ndarray, tuple[Any, …]]], …]

stacked_component_params(dists, engine)[source]

Return a generic stacked-scoring route for homogeneous components.

Generated declaration-backed routes are attempted first. Families with object lookups, table layouts, derived parameters, or wrapper structure can keep explicit backend_stacked_* hooks and still compose through this single dispatcher.

Parameters:
Return type:

StackedComponentParams

stacked_component_log_density(enc, route, engine)[source]

Evaluate a route returned by stacked_component_params.

Parameters:

  • enc (Any)

  • route (StackedComponentParams)

  • engine (ComputeEngine)

Return type:

Any

stacked_component_sufficient_statistics(enc, weights, route, engine, estimator=None)[source]

Return component-stacked legacy sufficient statistics for a stacked route.

Parameters:

  • enc (Any)

  • weights (Any)

  • route (StackedComponentParams)

  • engine (ComputeEngine)

  • estimator (ParameterEstimator | None)

Return type:

Any

unstack_component_stats(value, num_components)[source]

Return per-component legacy statistics from a component-stacked payload.

Parameters:

  • value (Any)

  • num_components (int)

Return type:

tuple[Any, …]

class StackedMixtureKernel(dist, engine, estimator=None)[source]

Bases: Kernel

Homogeneous mixture kernel with stacked component parameters.

The mixture mechanics live here: component matrix scoring, row-wise logsumexp, posterior weights, and legacy sufficient-stat dispatch. The leaf family still owns the actual component log-density math through backend_stacked_params and backend_stacked_log_density.

Parameters:

  • dist (Any)

  • engine (ComputeEngine)

  • estimator (ParameterEstimator | None)

component_type: type[Any]
encode(data)[source]

Encode raw observations with the mixture’s ordinary encoder.

Parameters:

data (Any)

Return type:

Any

component_scores(enc)[source]

Return unweighted component log densities with shape (n, k).

Parameters:

enc (Any)

Return type:

Any

score(enc)[source]

Return row log densities after adding mixture log weights.

Parameters:

enc (Any)

Return type:

Any

posteriors(enc)[source]

Return posterior component weights for each encoded row.

Parameters:

enc (Any)

Return type:

Any

property has_resident_accumulate: bool

Return true when the leaf family can accumulate sufficient stats on the engine.

resident_accumulate(enc, weights)[source]

Return engine-resident mixture sufficient statistics.

The mixture owns only the posterior mechanics. Leaf families own the sufficient-statistic algebra through backend_stacked_sufficient_statistics.

Parameters:
Return type:

StackedMixtureResidentStats

accumulate(enc, weights)[source]

Return mixture sufficient statistics in the legacy estimator format.

Parameters:
Return type:

Any

refresh(dist)[source]

Refresh stacked parameters after an M-step without changing structure.

Parameters:

dist (SequenceEncodableProbabilityDistribution)

Return type:

None

stacked_component_strategy(dist_type)[source]

Describe how homogeneous mixture component scoring will be dispatched.

Parameters:

dist_type (type[Any])

Return type:

str

class StackedMixtureKernelFactory(fallback=None)[source]

Bases: KernelFactory

Factory for homogeneous mixtures with distribution-owned stacked math.

Parameters:

fallback (KernelFactory | None)

build(dist, engine, estimator=None)[source]

Build a stacked mixture kernel when safe, otherwise use fallback.

Parameters:

  • dist (SequenceEncodableProbabilityDistribution)

  • engine (ComputeEngine)

  • estimator (ParameterEstimator | None)

Return type:

Kernel