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. Complex NumPy input is rejected when dtype is not given explicitly: the dtype inference below has no complex-precision policy of its own, and letting it fall through to the generic non-floating/non-Boolean case would silently discard the imaginary component and collapse the real part to an integer. Pass an explicit dtype=torch.complex64/torch.complex128 for a genuine complex tensor.

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 under the engine’s dtype and device policy.

Elements route through asarray() first. Leaf scorers legitimately hand back float64 NumPy arrays; stacking them raw produced a float64 tensor that later met float32 engine tensors – invisible on CPU, but a hard “mat1 double, mat2 float” error inside the HMM forward pass on the first real CUDA execution (rented RTX 3060, 2026-08-04). The engine’s dtype policy has to apply at every ingestion point, not only the explicit asarray calls.

Parameters:
Return type:

Any

concatenate(arrays, axis=0)[source]

Join tensors with torch.cat under the same ingestion policy as stack().

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).

matmul(x, y)[source]

Matrix-multiply under the engine’s dtype policy.

Both operands route through asarray(). Mixed float32/float64 operands reached this point through paths that never crossed an explicit ingestion call, and torch.matmul promotes on CPU but hard-errors on CUDA (“mat1 double, mat2 float”) – found and probed on the first real CUDA execution of the HMM forward pass (rented RTX 3060, 2026-08-04: mat1 float64 [100, 2] against mat2 float32 [2, 2]). The engine’s declared dtype IS the policy for floating operands, so applying it here is a correctness boundary, not a cast of convenience.

Parameters:
Return type:

Any

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