mixle.stats.compute.decomposition module¶

The model-parallel decomposition contract.

A distribution optionally declares how its parameters and sufficient statistics may be split across devices, via SequenceEncodableProbabilityDistribution.decomposition(). The descriptor names the axis the node splits along (its mixture components, its composite factors, its sequence/document units, …), the reduction that recombines per-shard sufficient statistics, whether the split is exact, and which children are shared (held whole and reduced, e.g. an HMM transition matrix or LDA topics).

This is the unifying primitive the structural planner and the model-parallel executor both consume. It is strictly opt-in: any node that does not override decomposition() reports Decomposition.atomic() (replicated, not split), so the contract never disturbs the existing data-parallel-with-replicated-model path. The cross-shard reduction is always the same additive combine() monoid the accumulators already implement – this introduces no new reduction algebra.

See ~/codex/notes/model-parallel-design.md (component C1) for the full design.

class DecompAxis(*values)[source]

Bases: Enum

The dimension a node splits along for model parallelism.

NONE = 'none'

COMPONENT = 'component'

FACTOR = 'factor'

SEQUENCE = 'sequence'

TOPIC = 'topic'

STATE = 'state'

class ReductionOp(*values)[source]

Bases: Enum

How per-shard sufficient statistics recombine at the cross-shard boundary.

SUM = 'sum'

LOGSUMEXP_RESPONSIBILITY = 'logsumexp'

REPLICATE = 'replicate'

class Decomposition(axis=DecompAxis.NONE, num_units=1, reduction=ReductionOp.REPLICATE, exact=True, child_roles=(), shared_children=(), engine_axis=None, key_pooling=False, extra=<factory>)[source]

Bases: object

How a distribution may be split across devices for model parallelism.

Parameters:

axis (DecompAxis)
num_units (int)
reduction (ReductionOp)
exact (bool)
child_roles (tuple[str, ...])
shared_children (tuple[str, ...])
engine_axis (int | None)
key_pooling (bool)
extra (dict[str, Any])

axis

the dimension the node splits along.

Type:: mixle.stats.compute.decomposition.DecompAxis

num_units

the number of splittable units along axis (components, factors, units).

Type:: int

reduction

how per-shard sufficient statistics recombine across shards.

Type:: mixle.stats.compute.decomposition.ReductionOp

exact

whether sharding this axis is exact EM (vs a restricted/approximate family).

Type:: bool

child_roles

names of the split children (e.g. ("component",)), informational.

Type:: tuple[str, …]

shared_children

children held whole on every shard and reduced, not split – e.g. an HMM transitions matrix or LDA topics (the structural statement of what cannot be split).

Type:: tuple[str, …]

engine_axis

the tensor axis for ComputeEngine.place_component_axis (DTensor sharding), or None when there is no homogeneous stacked-parameter tensor to shard.

Type:: int | None

key_pooling

whether per-shard estimates must route through keyed-tying (tie_component_shard_values) before the M-step.

Type:: bool

axis: DecompAxis = 'none'

num_units: int = 1

reduction: ReductionOp = 'replicate'

exact: bool = True

child_roles: tuple[str, ...] = ()

shared_children: tuple[str, ...] = ()

engine_axis: int | None = None

key_pooling: bool = False

extra: dict[str, Any]

classmethod atomic()[source]

The default: this node is not split – it is replicated across shards.

Return type:: Decomposition

property is_shardable: bool: True when this node declares a real (non-atomic) split axis with more than one unit.

register_decomposition(dist_type, decomposition)[source]

Parameters:

dist_type (type[Any])
decomposition (Decomposition)

Return type:

None

decomposition_for(x)[source]

Return the decomposition descriptor for a distribution instance or class.

Lookup mirrors mixle.stats.compute.capabilities.capabilities_for(): an instance decomposition() hook wins, then the class registry, then the MRO, then the atomic default.

Parameters:: x (Any)
Return type:: Decomposition