Compute Engines¶

Compute engines separate model semantics from array execution. A distribution owns the likelihood and sufficient-statistic math; an engine owns the array representation, precision, symbolic payload, or device boundary used to execute that math.

Use engines when you need GPU execution, JAX arrays, symbolic export, generated kernels, explicit precision control, or safe conversion between array backends.

Built-in Engines¶

NumpyEngine: Default host engine and the baseline for local kernels.
TorchEngine: Torch tensor engine for GPU/autograd-capable workflows and neural leaves.
JaxEngine: JAX array engine and bridge to JAX/NumPyro-oriented workflows.
SymbolicEngine: Symbolic expression engine for exporting log densities to SymPy, Sage, or LaTeX.

Basic Usage¶

Pass an engine to optimize without changing the model:

from mixle.engines import TorchEngine
from mixle.inference import optimize

model = optimize(
    data,
    estimator,
    engine=TorchEngine(device="cuda", dtype="float32"),
    out=None,
)

Move data folding to a backend separately:

model = optimize(data, estimator, backend="mp", num_workers=4, out=None)

engine= controls array math. backend= controls where encoded data are processed.

Torch DTensor Sharding¶

TorchEngine can represent component-sharded work through Torch DTensor when the installed Torch version supports the operations Mixle needs. In 0.6.2 the fully sharded component path is explicitly gated to Torch 2.5 or newer. Older Torch versions expose partial DTensor APIs but lack sharding strategies for operations used by mixture E-steps, which can otherwise fail deep inside Torch.

When the gate rejects a DTensor component-sharding request, use the engine-agnostic route instead:

model = optimize(data, estimator, backend="model_parallel", out=None)

The native model-parallel backend is the portable choice across Torch versions and devices. Use DTensor sharding only when the Torch runtime is new enough and you have a specific reason to keep component tensors resident in a distributed Torch mesh.

Engine Detection¶

engine_of detects the owning engine of an encoded payload. to_numpy is the explicit boundary for returning to NumPy.

from mixle.engines import engine_of, to_numpy

engine = engine_of(encoded_payload)
host_payload = to_numpy(encoded_payload)

Mixing incompatible array engines inside one payload raises an error instead of silently moving data across devices.

Precision¶

Precision helpers route computations explicitly:

from mixle.engines import auto_precision, engine_with_precision

precision = auto_precision(data, engine=engine)
engine = engine_with_precision(engine, precision)

optimize also accepts precision="auto" and precision="minimal". Use auto for device-aware defaults and minimal for data-aware reduced precision when verified safe.

The precision spectrum includes:

DoubleDouble, dd_sum, and dd_dot for extended precision;
Interval and sum_error_bound for error tracing;
AffineForm and allocate_precision for uncertainty-aware allocation;
FloatFormat, FixedPointFormat, and CodebookFormat for format experiments;
accurate_sum and sum_certificate for stable reductions.

Symbolic Export¶

Symbolic engines make density expressions inspectable:

from mixle.engines import SYMBOLIC_ENGINE, to_latex, to_sympy

symbolic = model.seq_log_density(encoded, engine=SYMBOLIC_ENGINE)
expr = to_sympy(symbolic)
latex = to_latex(symbolic)

Use this for reports, audits, or checking closed-form expressions. It is not intended to replace numeric fitting.

Registering Array Types¶

External array types can be associated with a compute engine:

from mixle.engines import register_array_type

register_array_type(MyArray, my_engine)

Registering an array type makes engine_of and recursive payload inspection route that type correctly.

Practical Guidance¶

Start with the default NumPy path until the model shape is correct.
Use TorchEngine for neural leaves and GPU-backed numeric work.
Use backend= for parallel or distributed data folding.
Use backend="model_parallel" for portable component parallelism across Torch versions.
Use symbolic export for inspection, not for production scoring.
Keep host/device boundaries explicit with to_numpy.
Use mixle.describe(model) to check whether a model supports backend scoring before assuming an engine will accelerate it.

API Map¶

Import	Purpose
`NumpyEngine`, `NUMPY_ENGINE`, `FUSED_NUMPY_ENGINE`	local host execution
`TorchEngine`	Torch tensors, GPU, autograd-aware workflows
`JaxEngine`	JAX arrays and JAX-oriented routes
`SymbolicEngine`, `SYMBOLIC_ENGINE`	symbolic payloads and expression export
`SymbolicExpression`	symbolic payload node used by symbolic export
`engine_of`, `to_numpy`, `register_array_type`	engine detection and explicit conversion
`normalize_numpy_dtype`, `normalize_torch_dtype`	dtype normalization for engine setup
`auto_precision`, `engine_with_precision`, `precision_name`	precision routing
`to_sympy`, `to_sage`, `to_latex`	symbolic export formats
`DoubleDouble`, `Interval`, `AffineForm`	precision and error-analysis tools
`float64_sum_is_accurate`	quick check for whether ordinary float64 summation is sufficient