mixle.engines.base module¶

Array compute-engine protocol used by backend-neutral kernels.

active_engine()[source]

Return the compute engine driving the current EM step, or None outside one.

Return type:: ComputeEngine | None

using_active_engine(engine)[source]

Mark engine active for the duration of the block (used by the estimation loop).

Parameters:: engine (Any)

class ComputeEngine[source]

Bases: ABC

Small array-backend interface for numpy/torch/etc.

Engines own arithmetic policy: array library, device, dtype, and optional compilation. Distribution and kernel code should depend only on this surface when it wants backend-neutral arrays.

name = 'base'

supports_autograd = False

dtype = None

device = 'cpu'

REQUIRED_OPS: tuple[str, ...] = ('asarray', 'zeros', 'empty', 'arange', 'to_numpy', 'stack', 'log', 'exp', 'sqrt', 'abs', 'where', 'maximum', 'clip', 'floor', 'isnan', 'isinf', 'sum', 'max', 'dot', 'matmul', 'cumsum', 'logsumexp', 'bincount', 'unique', 'searchsorted', 'index_add', 'gammaln', 'digamma', 'betaln', 'erf')

pi = 3.141592653589793

e = 2.718281828459045

euler_gamma = 0.5772156649015329

inf = inf

zero = 0.0

one = 1.0

two = 2.0

half = 0.5

constant(value)[source]

Return value in this engine’s scalar representation (identity for numeric engines).

Parameters:: value (Any)
Return type:: Any

supports_numba = False

resident_estep = True

property accumulator_dtype: Any

High-precision dtype for sufficient-statistic reductions, or None when not applicable.

Numeric engines override this with their float64 accumulator so a reduced-precision fit does not drift on large N (see NumpyEngine/TorchEngine). The base returns None – meaning “no separate accumulator dtype”, which is the correct policy for engines that never drive the numeric accumulate path (e.g. the symbolic engine, where reductions are exact expression trees). None is also a valid dtype= argument to sum (NumPy’s default).

property precision: str: Return the engine dtype policy as a stable user-facing name.

with_precision(precision)[source]

Return an equivalent engine with a different floating-point policy.

Parameters:: precision (Any)
Return type:: ComputeEngine

abstractmethod asarray(x, dtype=None)[source]

Convert x into this engine’s array/tensor representation.

Parameters:

x (Any)
dtype (Any)

Return type:

Any

abstractmethod zeros(shape, dtype=None)[source]

Allocate a zero-filled array on this engine.

Parameters:

shape (Any)
dtype (Any)

Return type:

Any

abstractmethod empty(shape, dtype=None)[source]

Allocate an uninitialized array on this engine.

Parameters:

shape (Any)
dtype (Any)

Return type:

Any

abstractmethod arange(*args, **kwargs)[source]

Return an evenly spaced one-dimensional array on this engine.

Parameters:

args (Any)
kwargs (Any)

Return type:

Any

abstractmethod to_numpy(x)[source]

Move an engine array back to a NumPy/host representation.

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

abstractmethod stack(arrays, axis=0)[source]

Stack a sequence of arrays along axis.

Parameters:

arrays (Any)
axis (int)

Return type:

Any

requires_grad(x)[source]

Return whether x participates in this engine’s autograd graph.

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

compile(fn)[source]

Optionally compile fn; engines without a compiler return it unchanged.

Parameters:: fn (Callable)
Return type:: Callable

replicate(x)[source]

Return x in the engine’s replicated placement, when applicable.

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

place_component_axis(x, axis=0)[source]

Return x with a component-axis placement, when the engine supports it.

Parameters:

x (Any)
axis (int)

Return type:

Any