mixle.engines.torch_engine module

Torch implementation of the ComputeEngine protocol.

class TorchEngine(device=None, dtype=None, compile=False, mesh=None, shard=None)[source]

Bases: ComputeEngine

Torch tensor engine for device placement, autograd, and optional DTensor sharding.

Parameters:

  • device (str | None)

  • dtype (Any)

  • compile (bool)

  • mesh (Any)

  • shard (str | None)

name = 'torch'
supports_autograd = True
property accumulator_dtype: Any

High-precision dtype for sufficient-statistic reductions (float64, or float32 on MPS).

Reductions that aggregate over observations accumulate in float64 even when scoring runs in reduced precision, so a float32 fit does not drift on large N. MPS has no float64, so there the accumulator falls back to float32 (its max precision) — fits on very large N may drift slightly.

with_precision(precision)[source]

Return a Torch engine with the same placement and a new dtype policy.

Parameters:

precision (Any)

Return type:

TorchEngine

asarray(x, dtype=None)[source]

Convert x to a Torch tensor or DTensor on the configured device.

Contract: float inputs are force-cast to the engine’s float dtype (e.g. a float32 numpy array becomes the engine’s float64) unless dtype is given; this differs from numpy’s asarray, which preserves the input dtype.

Parameters:
Return type:

Any

zeros(shape, dtype=None)[source]

Allocate a zero tensor with this engine’s dtype/device/placement.

Parameters:
Return type:

Any

empty(shape, dtype=None)[source]

Allocate an uninitialized tensor with this engine’s dtype/device/placement.

Parameters:
Return type:

Any

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

Return torch.arange on the configured device.

Parameters:
Return type:

Any

to_numpy(x)[source]

Gather/detach a Torch tensor or DTensor to a host NumPy array.

Parameters:

x (Any)

Return type:

ndarray

stack(arrays, axis=0)[source]

Stack tensors with torch.stack and apply default placement.

Parameters:
Return type:

Any

replicate(x)[source]

Return x replicated across the configured DeviceMesh.

Parameters:

x (Any)

Return type:

Any

place_component_axis(x, axis=0)[source]

Place a stacked component parameter tensor on Shard(axis).

TorchEngine(mesh=..., shard='components') is the model-parallel configuration from the compute-engine design. Encoded data remains replicated, while homogeneous stacked-kernel parameter tensors are sharded along their component dimension. With no mesh, or with no component sharding requested, this is just asarray.

Parameters:
Return type:

Any

requires_grad(x)[source]

Return whether x is a Torch tensor/DTensor requiring gradients.

Parameters:

x (Any)

Return type:

bool

compile(fn)[source]

Compile fn with torch.compile when enabled and available.

Parameters:

fn (Callable)

Return type:

Callable

static log(x)
static exp(x)
static sqrt(x)
static abs(x)
static where(*args)
static maximum(x, y)
static clip(x, a_min=None, a_max=None)
static floor(x)
static isnan(x)
static isinf(x)
static sum(x, *args, **kwargs)[source]

Return torch.sum accepting either axis or dim.

static max(x, *args, **kwargs)[source]

Return torch.max accepting either axis or dim (incl. a tuple of axes).

static dot(x, y)
static matmul(x, y)
static cumsum(x, *args, **kwargs)[source]

Return torch.cumsum accepting axis/dim and defaulting to a flattened scan.

static logsumexp(x, *args, **kwargs)[source]

Return torch.logsumexp accepting either axis or dim.

torch < 2.5 registers no DTensor sharding strategy for logsumexp, so reducing over a sharded axis (the mixture E-step’s log-partition over component-sharded scores) raises NotImplementedError. Fall back to redistributing the DTensor to replicated first – the reduction then runs locally and is correct on any torch version. torch >= 2.5 has the strategy and takes the fast native path unchanged; non-DTensor inputs re-raise the original error.

static bincount(x, *args, **kwargs)
static unique(x, *args, **kwargs)
static searchsorted(x, y, *args, **kwargs)
static gammaln(x)
static digamma(x)
static betaln(x, y)
static erf(x)
static cos(x)
static sin(x)
static arctan2(x, y)
static i0e(x)
static erfcx(x)
hmm_forward_ll(log_emit, log_w, log_a, mask)[source]

Fused log-space forward: per-sequence emission log-likelihood (freeze-alpha padding).

Parameters:
Return type:

Any

hmm_forward_backward(log_emit, log_w, log_a, mask, weights=None)[source]

Fused log-space Baum-Welch recurrences, exactly mirroring the generic engine-op version: alpha freezes across padded steps, beta carries, gamma is mask-SELECTED (NaN-safe for empty sequences), xi is computed for all transitions at once, and per-sequence weights scale gamma / xi / pi. Returns (ll, gamma, xi_sum, pi).

Parameters:
Return type:

tuple[Any, Any, Any, Any]

index_add(out, index, values)[source]

Add values into out along axis 0 using Torch index_add.

Contract: return-value-only – torch.Tensor.index_add (no trailing underscore) does not mutate out in place, unlike numpy’s add.at; callers must use the returned tensor.

Parameters:
Return type:

Any