mixle.inference.backends module

Inference-backend registry — register, don’t branch for mixle.inference.

Mirrors mixle.stats.compute.kernel.register_kernel_factory(): each engine’s NUTS implementation self-registers an InferenceBackend at import time, so the dispatcher (mixle.inference.nuts()) never grows a central if engine == ... switch. A backend declares

  • name — the selector string ("numpy", "numba", "torch", "jax").

  • available — a zero-arg predicate: is the engine importable on this host? Kept lazy so import mixle.inference works with any subset of optional engines installed.

  • target_kind — what contract the caller’s target must satisfy: a numpy fused value_and_grad ("numpy_vg"), an @njit fused value_and_grad ("njit_vg"), a torch scalar logp ("torch_logp"), or a jax scalar logp ("jax_logp"). The kinds cannot be auto-converted across autodiff systems, so the target is what ultimately picks a backend in "auto" mode (see select_backend()).

  • nuts — the callable that runs the sampler and returns a mixle.inference.NutsResult.

available_backends() lists the installed engines; select_backend() resolves the backend= argument (including "auto").

class InferenceBackend(name, available, target_kind, nuts)[source]

Bases: object

A registered inference engine: a name, an availability probe, a target contract, a sampler.

Parameters:

  • name (str)

  • available (Callable[[], bool])

  • target_kind (str)

  • nuts (Callable[..., NutsResult])

name: str
available: Callable[[], bool]
target_kind: str
nuts: Callable[..., NutsResult]
register_inference_backend(backend)[source]

Register (or replace) an inference backend under backend.name.

Parameters:

backend (InferenceBackend)

Return type:

None

get_inference_backend(name)[source]

Return the registered backend named name (raises if unknown).

Parameters:

name (str)

Return type:

InferenceBackend

available_backends()[source]

Return the names of registered backends whose engine is importable, in registration order.

Return type:

list[str]

select_backend(backend='auto', target=None)[source]

Resolve a backend= argument to a concrete, available backend name.

Policy:

  • An explicit backend (anything but "auto") is honored — it must be registered and its engine importable, else a clear error.

  • "auto" with a target kind hint picks the first preferred-and-available backend for that kind (e.g. "torch_logp" -> torch; "numpy_vg" -> numpy, then numba). This keeps the always-available numpy path the default for plain numpy targets.

  • "auto" with no hint falls back to the first available backend, preferring "numpy" (the always-present, dependency-free path) when it is available.

Raises if nothing is available or the explicit choice is unavailable.

Parameters:

  • backend (str)

  • target (str | None)

Return type:

str