mixle.engines.torch_engine module

Torch implementation of the ComputeEngine protocol.

The engine handles tensor placement, dtype policy, autograd support, optional compilation, and component sharding for resident scoring and estimation paths.

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)

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

sum(x, *args, **kwargs)[source]

Return torch.sum accepting either axis or dim.

Promotes a floating input to accumulator_dtype when the caller doesn’t pass an explicit dtype= – mirroring NumpyEngine.sum(). Without this, torch.sum accumulates in the input tensor’s own dtype by default, so a float32-precision fit on this engine would silently drift on large N (the exact catastrophic-cancellation risk accumulator_dtype exists to guard against) while the numpy engine, which already promotes, stayed accurate.

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

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

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.

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