Nx.Vulkan.Fast (nx_vulkan v0.2.0)

Copy Markdown View Source

Named fused kernels for MCMC hot paths.

Each function is a composition of standard Nx ops that produces a mathematically-equivalent result to the fused shader that Nx.Vulkan.VulkanoBackend would dispatch. Cross-backend correctness is guaranteed (EXLA, BinaryBackend, VulkanoBackend).

Note on Nx 0.12 migration

Prior to Nx 0.12, each function emitted Nx.Defn.Expr.optional/3 IR nodes (an internal Nx API, so not linked here) for backend-specific fused dispatch. That API was removed in Nx 0.12. The functions now call the fallback Nx ops directly. The VulkanoBackend's per-op dispatch is fast enough that the fused-kernel optimization is not critical — the chain shader path (where performance matters) bypasses this module entirely.

How to use

Inside a defn or any Nx.Defn.jit-traced function:

defn leapfrog_step(q, eps, p, grad) do
  q_new = Nx.Vulkan.Fast.leapfrog_position(q, eps, p)
  p_new = Nx.Vulkan.Fast.momentum_step(p, eps, grad)
  {q_new, p_new}
end

Summary

Functions

inv_mass_apply(p, inv_mass)

Apply diagonal mass-matrix inverse: p * inv_mass. Trivial as a fused kernel (one binary op), but named for symmetry — a future shader could combine it with adjacent ops in the leapfrog without changing call sites.

kinetic_energy(p, inv_mass)

Kinetic energy: 0.5 * sum(p² * inv_mass). Reduces to a scalar. Used in NUTS for the joint log-probability: joint_logp = log_prob - kinetic_energy(p, inv_mass).

leapfrog_momentum_half(p, half_eps, grad)

Half-step momentum update: p + half_eps * grad. Used at the start and end of every leapfrog iteration in the standard symplectic integrator. half_eps is eps / 2 precomputed by the caller.

leapfrog_position(q, eps, p)

Position update: q + eps * p. The dominant elementwise body in every NUTS leapfrog.

momentum_step(p, eps, grad)

Full-step momentum update: p + eps * grad. Same shape as the half-step but kept distinct to signal the caller's intent.

normal_logpdf(x, mu, sigma)

Normal log-density: -0.5*((x-mu)/sigma)² - log(sigma) - 0.5*log(2π). Output shape matches x. The MCMC distribution-density hot path.

Functions

inv_mass_apply(p, inv_mass)

@spec inv_mass_apply(Nx.t(), Nx.t()) :: Nx.t()

Apply diagonal mass-matrix inverse: p * inv_mass. Trivial as a fused kernel (one binary op), but named for symmetry — a future shader could combine it with adjacent ops in the leapfrog without changing call sites.

kinetic_energy(p, inv_mass)

@spec kinetic_energy(Nx.t(), Nx.t()) :: Nx.t()

Kinetic energy: 0.5 * sum(p² * inv_mass). Reduces to a scalar. Used in NUTS for the joint log-probability: joint_logp = log_prob - kinetic_energy(p, inv_mass).

leapfrog_momentum_half(p, half_eps, grad)

@spec leapfrog_momentum_half(Nx.t(), Nx.t(), Nx.t()) :: Nx.t()

Half-step momentum update: p + half_eps * grad. Used at the start and end of every leapfrog iteration in the standard symplectic integrator. half_eps is eps / 2 precomputed by the caller.

leapfrog_position(q, eps, p)

@spec leapfrog_position(Nx.t(), Nx.t(), Nx.t()) :: Nx.t()

Position update: q + eps * p. The dominant elementwise body in every NUTS leapfrog.

Examples

iex> q = Nx.tensor([1.0, 2.0])
iex> eps = Nx.tensor([0.5, 0.5])
iex> p = Nx.tensor([2.0, 4.0])
iex> Nx.Vulkan.Fast.leapfrog_position(q, eps, p) |> Nx.to_flat_list()
[2.0, 4.0]

momentum_step(p, eps, grad)

@spec momentum_step(Nx.t(), Nx.t(), Nx.t()) :: Nx.t()

Full-step momentum update: p + eps * grad. Same shape as the half-step but kept distinct to signal the caller's intent.

normal_logpdf(x, mu, sigma)

@spec normal_logpdf(Nx.t(), Nx.t(), Nx.t()) :: Nx.t()

Normal log-density: -0.5*((x-mu)/sigma)² - log(sigma) - 0.5*log(2π). Output shape matches x. The MCMC distribution-density hot path.