mixle.inference.mcmc.gradients module

Autograd gradients for MCMC log-targets (optional Torch backend).

Hamiltonian samplers (HMC, MALA, NUTS) need the gradient of the log-target. The default finite-difference gradient costs O(d) full target evaluations per step – the documented bottleneck for parameter-posterior HMC over higher-dimensional models. When the target is expressed in Torch, autograd returns the exact gradient in a single backward pass, independent of dimension. This module provides that bridge.

It is entirely optional: import-light, and torch_gradient() raises a clear error only if called without Torch. Callers that want a guaranteed-available gradient can fall back to the finite-difference path in parameter_bridge.

torch_available()[source]

Return True if Torch can be imported (autograd gradients are available).

Return type:

bool

torch_gradient(log_target_torch, dtype='float64')[source]

Build an exact grad_log_target from a Torch-valued log-target via autograd.

Parameters:
  • log_target_torch (Callable[[Any], Any]) – Callable mapping a 1-D Torch tensor of parameters to a scalar Torch tensor (the unnormalized log target). Differentiation flows through whatever Torch ops it uses, so the gradient is exact – one backward pass regardless of dimension.

  • dtype (str) – Torch float dtype for the parameter tensor (“float64” by default for MCMC numerics; use “float32” to match a single-precision model).

Returns:

grad(x) -> gradient accepting and returning numpy (a float for scalar x, otherwise an array shaped like x). Suitable as the grad_log_target argument to hamiltonian_monte_carlo(), nuts(), or a Langevin proposal.

Return type:

Callable[[Any], Any]

value_and_torch_gradient(log_target_torch, dtype='float64')[source]

Like torch_gradient() but returns (value, gradient) in one backward pass.

Useful for samplers that need both the log-target and its gradient at the same point (HMC/NUTS leapfrog), avoiding a redundant forward evaluation.

Parameters:
Return type:

Callable[[Any], tuple[float, Any]]