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

register_inference_backend(backend, *, replace=False)[source]

Register an inference backend under backend.name.

The registry is global process state that select_backend() dispatches every nuts call through, so one malformed registration is not a local mistake – it corrupts dispatch for every caller. A backend with an empty name and an unrecognized target_kind used to register successfully and become selectable, even though select_backend() rejects that same kind string when a caller supplies it.

Replacing an already-registered name requires replace=True. Silent replacement let a second import or a typo’d name swap out a live sampler with no trace; an intentional override (a test double, a vendored engine) says so explicitly.

Raises:

  • TypeError – if backend is not an InferenceBackend, or available/nuts are not callable.

  • ValueError – for an empty/blank name, an unrecognized target_kind, or an unrequested replacement of a registered name.

Parameters:

  • backend (InferenceBackend)

  • replace (bool)

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.

Each probe is isolated: one backend whose available() raises used to propagate out of this call and hide every healthy backend, so a single broken third-party registration made the always-present numpy path look unavailable too. A raising probe answers “unknown”, which is reported as a warning and treated as unavailable – it cannot be silently equivalent to a probe that returned False, because the caller may need to know its engine is misconfigured rather than absent.

Return type:

list[str]

backend_availability()[source]

Per-backend (name, available, probe_error) in registration order.

probe_error is None unless that backend’s available() raised, in which case it is the formatted exception and available is False.

Return type:

list[tuple[str, bool, str | None]]

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. A declared kind is a contract, not a preference: if no backend advertising that contract is available, this raises rather than silently handing e.g. a torch/jax scalar logp to a backend that will call it expecting a numpy value_and_grad -> (value, gradient) pair. The kinds are not convertible across autodiff systems, so there is no correct fallback between them.

  • "auto" with no hint (target=None) falls back to the first available backend, preferring "numpy" (the always-present, dependency-free path) when it is available. This is the genuinely undetermined case – a numpy and an @njit value_and_grad are indistinguishable plain callables – and stays unchanged.

Raises:

  • ValueError – if target is not one of the known target kinds.

  • RuntimeError – if nothing is available, the explicit choice is unavailable, or no available backend implements the declared target kind.

Parameters:

  • backend (str)

  • target (str | None)

Return type:

str