API Reference

needs_host_staging(x::AbstractArray) -> Bool

Whether x has to be filled on the host and copied over, rather than written element by element in place. false for anything with cheap scalar indexing, which is the default.

ReactantExt also reads it as "this array lives on a device", to warn when the XLA client is on the host while the parameters are not.

CUDAExt defines it for CuArray. To support another device array type:

ParallelMCMC.needs_host_staging(::ROCArray) = true

DensityModel(logdensity, grad_logdensity, dim; param_names, logdensity_batch, grad_logdensity_batch, hvp, hvp_batch)

Wrap a log-density and its derivatives for ParallelMCMC samplers.

Each derivative slot (grad_logdensity, hvp, grad_logdensity_batch, hvp_batch) takes a callable or an ADTypes.AbstractADType, so a model can be built from the log-density alone:

DensityModel(logp, AutoForwardDiff(), dim)

Backends are prepared when sampling starts. An HVP backend differentiates a callable gradient once. With an AD-derived gradient it forms DifferentiationInterface.SecondOrder(hvp_backend, grad_backend) on the log-density. An explicit SecondOrder also bypasses the gradient slot.

• logdensity(x::AbstractVector) -> Real
• grad_logdensity — callable x -> AbstractVector, or an AD backend.
• hvp — optional callable (x, v) -> AbstractVector, or a backend. If nothing, DEER builds the HVP from the sampler's backend.
• logdensity_batch(X::AbstractMatrix) -> AbstractVector — optional batched log-density over independent columns.
• grad_logdensity_batch — optional callable X -> AbstractMatrix, or a backend. It is derived from logdensity_batch when grad_logdensity is a backend.
• hvp_batch — optional callable (X, V) -> AbstractMatrix, or a backend, resolved like hvp.
• dim::Int — dimensionality of the parameter space
• param_names — optional collection of parameter names used in FlexiChains output. If nothing (the default), uses a single vector-valued parameter :x with shape (dim,). See the Parameter names section of the docs for more information.

Both batched derivative slots require logdensity_batch. AutoReactant() has additional pairing and tracing constraints; see the GPU guide.

MALASampler(epsilon; cholM=nothing)

Metropolis-Adjusted Langevin Algorithm sampler with step size epsilon.

Optionally pass cholM = cholesky(M) to use a mass matrix M as a preconditioner. The proposal becomes y = x + ε M ∇logp(x) + √(2ε) L ξ where L is the Cholesky factor of M.

AdaptiveMALASampler(epsilon_init; n_warmup, target_accept, gamma, t0, kappa, cholM)

ParallelMALASampler(epsilon; T, maxiter, tol_abs, tol_rel, jacobian, damping, probes, cholM, backend)

DEER-parallelized MALA sampler.

Supported Jacobian modes are :stoch_diag (the default Hutchinson diagonal estimator) and :diag (exact diagonal via D JVPs).

backend supplies Hessian-vector products when the DensityModel brings no hvp / hvp_batch of its own. It does not supply gradients. Leave it out when the model carries its own HVPs.

With backend=ADTypes.AutoReactant(), grad_logdensity must be AutoReactant() or a callable.

callback and initial_state route through the generic step-wise loop instead.

MALATapeElement(ξ, u)

One element of the MALA noise tape: a noise vector ξ ~ N(0,I) and a uniform scalar u ~ Uniform(0,1). Stored with a concrete vector type V for type stability inside DEER.TapedRecursion.

State for a MALASampler chain. Holds the prepped model so AD preparation happens once per chain.

One MALASampler sample: parameter vector x, its log-density logp, and an accept/reject flag.

State for an AdaptiveMALASampler chain, including dual-averaging adaptation statistics. Holds the prepped model so AD preparation happens once per chain.

One AdaptiveMALASampler sample: parameter vector x, its log-density logp, the step size used for that transition, and whether the sample came from warmup.

State for a ParallelMALASampler chain. Holds the prepped model so AD preparation happens once per chain.

One ParallelMALASampler sample: parameter vector x and its log-density logp.

Deterministic recursion driven by a pre-generated tape.

• step_fwd(x, tape_t): exact forward transition (may include MH accept/reject).
• jvp(x, tape_t, v): explicit JVP (x, tape_t, v) -> J·v of the differentiable surrogate step. Required; must not call AD through step_fwd's control flow.
• fwd_and_jvp(x, tape_t, v): optional fused function -> (x_next, J·v).
• fwd_and_jvp_batch(Xbar, Z): optional batched function -> (FT, Jt).

Reusable buffers for allocation-light DEER updates and solves.

All buffers are created with the same array type / device placement as S_template (or s0_template for vector buffers), so the workspace is GPU-compatible when constructed from device-array templates.

Run DEER iterations until convergence.

When no workspace is supplied, one is created automatically. Supplying a pre-allocated DEERWorkspace is the intended path for repeated GPU solves.

If workspace is supplied, copy_result defaults to true so the returned trajectory is owned by the caller and will not be overwritten by a later solve using the same workspace. Set copy_result=false for allocation-sensitive internal loops; in that mode the returned trajectory may alias workspace-owned buffers such as workspace.S_tmp.

Reusable workspace for the diagonal affine scan.

Buffers are allocated with the same array type / device placement as the template matrix used to construct the workspace, so this works on GPU arrays and CPU arrays alike.

Stop-gradient surrogate step used for Jacobians.