mixle.utils.parallel.model_parallel module

Model-parallel estimation: distribute a model’s shardable axis across workers.

The inversion of the data-parallel backends: there the model is replicated and the data sharded. Here the model’s shardable axis is distributed. Two entry points share one recursive fold:

• optimize(..., backend="model_parallel") – the ModelParallelEncodedData handle: the data is replicated in-process and the model axis is threaded. Simplest, single-machine.

• optimize(ModelParallelEstimator(est), backend="spark"|"mpi"|"mp"|"local") – the estimator wrapper: composes with any data backend, so the data is sharded by that backend (Spark partitions, MPI ranks, mp pool) while the model axis is distributed inside each partition’s accumulator. This is the data x model composition – both axes at once.

The fold (model_parallel_fold()) is recursive: it walks the whole model tree and threads the axes that a per-node compute-cost model says save the most wall-time (_parallel_ids, consistent with the structural planner – a narrow batch of heavy MVGaussians beats a wider batch of low-cost leaves), recursing serially below any threaded node so no two pools ever nest. Each recursive case reproduces the corresponding accumulator’s seq_update exactly:

• FACTOR (Composite/Record) – the per-factor accumulators are independent, so the per-factor seq_update calls are distributed (bit-identical).

• COMPONENT (mixtures) – the responsibility logsumexp couples the components, so the low-cost normalization runs centrally on the gathered score matrix while the expensive per-component scoring and accumulation are distributed – a bit-identical mirror of MixtureAccumulator.seq_update.

• atomic / unknown – the replicated base case acc.seq_update(enc, weights, model).

So the whole fold is bit-identical to the single-node path (the data-axis reduce across partitions is the usual additive combine, exact up to float reassociation like every data-parallel backend). Correct for every family; never worse than backend="local".

class ModelParallelEncodedData(data, *, estimator=None, model=None, encoder=None, num_workers=None, **_)

Bases: EncodedDataHandle

Replicate the data, distribute the model’s shardable axis across threads (single machine).

Parameters:

• data (Any)

• estimator (Any | None)

• model (Any | None)

• encoder (Any | None)

• num_workers (int | None)

• _ (Any)

pysp_seq_log_density_sum(estimate)

Return the encoded-data size and total log likelihood under estimate.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)

Run one model-parallel E/M update from prev_estimate.

Parameters:

• estimator (Any)

• prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)

Initialize a model by randomly selecting observations with probability p.

Parameters:

• estimator (Any)

• rng (RandomState)

• p (float)

Return type:

Any

pysp_stream_accumulate(estimator, model)

Accumulate model-parallel sufficient statistics for streaming backends.

Parameters:

• estimator (Any)

• model (Any)

Return type:

tuple[float, Any]

class ModelParallelEstimator(inner, num_workers=None)

Bases: ParameterEstimator

Wrap an estimator so EM distributes the model axis – composing with any backend= for the data.

optimize(ModelParallelEstimator(est), backend="spark"|"mpi"|"mp"|"local") shards the data through that backend while each partition’s E-step distributes the model axis. The M-step (estimate) and the accumulator’s value/combine contract are the wrapped estimator’s, unchanged.

Parameters:

• inner (ParameterEstimator)

• num_workers (int | None)

accumulator_factory()

Return a factory that wraps the inner estimator’s accumulator factory.

Return type:

ModelParallelAccumulatorFactory

estimate(nobs, suff_stat)

Delegate the M-step to the wrapped estimator.

Parameters:

• nobs (float | None)

• suff_stat (Any)

Return type:

Any

class ModelParallelAccumulator(inner, num_workers=None)

Bases: SequenceEncodableStatisticAccumulator

Wrap an accumulator so its E-step seq_update runs the recursive model-parallel fold.

All sufficient-statistic methods delegate to the wrapped (inner) accumulator unchanged, so the value/combine/from_value/key-merge contract – and thus every data backend’s reduce – is preserved; only seq_update is replaced with the distributed fold. Holding inner in vars() keeps the accumulator key-validator’s recursion transparent.

Parameters:

• inner (SequenceEncodableStatisticAccumulator)

• num_workers (int | None)

update(x, weight, estimate)

Delegate scalar accumulation to the wrapped accumulator.

Parameters:

• x (Any)

• weight (float)

• estimate (Any)

Return type:

None

initialize(x, weight, rng)

Delegate scalar initialization to the wrapped accumulator.

Parameters:

• x (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_update(x, weights, estimate)

Run the wrapped accumulator’s sequence update through the model-parallel fold.

Parameters:

• x (Any)

• weights (ndarray)

• estimate (Any)

Return type:

None

seq_initialize(x, weights, rng)

Delegate encoded initialization to the wrapped accumulator.

Parameters:

• x (Any)

• weights (ndarray)

• rng (Any)

Return type:

None

combine(suff_stat)

Merge sufficient statistics into the wrapped accumulator.

Parameters:

suff_stat (Any)

Return type:

ModelParallelAccumulator

value()

Return the wrapped accumulator’s sufficient-statistic value.

Return type:

Any

from_value(x)

Replace the wrapped accumulator from a sufficient-statistic value.

Parameters:

x (Any)

Return type:

ModelParallelAccumulator

key_merge(stats_dict)

Delegate keyed statistic merging to the wrapped accumulator.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)

Delegate keyed statistic replacement to the wrapped accumulator.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()

Return the wrapped accumulator’s compatible data encoder.

Return type:

Any

class ModelParallelAccumulatorFactory(inner_factory, num_workers=None)

Bases: StatisticAccumulatorFactory

Factory that wraps another accumulator factory with model-parallel updates.

Parameters:

• inner_factory (Any)

• num_workers (int | None)

make()

Create a fresh model-parallel accumulator wrapper.

Return type:

ModelParallelAccumulator

model_parallel_fold(acc, model, enc, weights, num_workers=None)

Run a model-parallel E-step of model over encoded enc into accumulator acc (in place).

Parameters:

• acc (Any)

• model (Any)

• enc (Any)

• weights (ndarray)

• num_workers (int | None)

Return type:

None

auto_parallel_estimator(estimator, model, resources=None, *, n_data=None, min_components_per_shard=1)

Consult the C2 planner (decompose_model()) and return (estimator, decomposition).

When the planner picks model-parallelism for model on resources, the estimator is wrapped in ModelParallelEstimator sized to the planner’s cuts; otherwise the plain estimator is returned (replicate the model, shard the data – already optimal when N dominates). Either way, run the returned estimator through optimize(data, est, backend=<data backend>): the data axis is handled by backend and the model axis, if any, by the wrapper – composing into the data x model split. The decomposition’s rationale explains the choice. resources defaults to the local CPU slots (use Resources.from_spark(sc) / Resources.from_mpi() to size the split to a real cluster).

Parameters:

• estimator (Any)

• model (Any)

• resources (Any)

• n_data (int | None)

• min_components_per_shard (int)

Return type:

tuple[Any, Any]