mlxmc 0.1.0__py3-none-any.whl

mlxmc/__init__.py

@@ -0,0 +1,26 @@
+"""mlxmc — MCMC samplers in Apple MLX.
+
+Affine-invariant ensemble, HMC (identity / preconditioned), Stan-style warmup
+adaptation, and NUTS, plus ESS diagnostics and a set of example targets. Every
+sampler takes a single-point log-density `logp(x) -> scalar` for `x` of shape
+`(D,)`; batching over chains/walkers is handled internally with `vmap`.
+"""
+from mlxmc import targets
+from mlxmc.diagnostics import autocorr_1d, integrated_time, report
+from mlxmc.ensemble import make_sampler, run_ensemble
+from mlxmc.hmc import make_hmc, run_hmc
+from mlxmc.nuts import make_nuts, nuts_warmup, run_nuts
+from mlxmc.preconditioned import make_phmc, run_phmc
+from mlxmc.warmup import DualAveraging, make_warmup_step, run_chain, warmup
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "make_sampler", "run_ensemble",
+    "make_hmc", "run_hmc",
+    "make_phmc", "run_phmc",
+    "DualAveraging", "make_warmup_step", "warmup", "run_chain",
+    "make_nuts", "run_nuts", "nuts_warmup",
+    "autocorr_1d", "integrated_time", "report",
+    "targets",
+]

mlxmc/diagnostics.py

@@ -0,0 +1,49 @@
+"""Effective sample size: how many *independent* draws the chains are worth.
+
+ESS = N / tau, with tau the integrated autocorrelation time (emcee-style:
+FFT autocorrelation averaged over walkers, Sokal automatic windowing). The
+fair cross-sampler metric is ESS/sec, since it folds in per-step cost.
+
+Pure numpy on a structured (T, N, D) chain -- no sampler or MLX dependency, so
+it's a leaf module the samplers don't pull in.
+"""
+import numpy as np
+
+
+def autocorr_1d(x):
+    x = x - x.mean()
+    n = len(x)
+    f = np.fft.fft(x, n=2 * n)
+    acf = np.fft.ifft(f * np.conj(f))[:n].real
+    if acf[0] == 0:               # constant (stuck) walker: autocorrelation undefined
+        return None
+    return acf / acf[0]
+
+
+def auto_window(taus, c=5.0):
+    m = np.arange(len(taus)) < c * taus
+    return np.argmin(m) if np.any(~m) else len(taus) - 1
+
+
+def integrated_time(y):                       # y: (T, N) for one dimension
+    # Skip stuck (zero-variance) walkers, which would make acf/acf[0] a 0/0 = nan and
+    # poison the walker average. A large skipped fraction is itself a mixing-failure tell.
+    acfs = [a for a in (autocorr_1d(y[:, w]) for w in range(y.shape[1])) if a is not None]
+    if not acfs:
+        return np.nan
+    f = np.mean(acfs, axis=0)
+    taus = 2.0 * np.cumsum(f) - 1.0
+    return taus[auto_window(taus)]
+
+
+def report(chain_mx, label, dt):
+    c = np.array(chain_mx)                     # (T, N, D)
+    T, N, D = c.shape
+    total = T * N
+    tau = max(integrated_time(c[:, :, d]) for d in range(D))
+    ess = total / tau
+    print(f"\n[{label}]")
+    print(f"  raw samples {total:,}  ({T} steps x {N})   wall {dt:.2f}s")
+    print(f"  tau (worst dim): {tau:.1f} steps  ->  ESS {ess:,.0f}  ({100 * ess / total:.1f}% independent)")
+    print(f"  ESS/sec: {ess / dt:,.0f}")
+    return ess, dt

mlxmc/ensemble.py

@@ -0,0 +1,60 @@
+"""Affine-invariant ensemble sampler (Goodman & Weare 2010) in MLX.
+
+This is the algorithm behind `emcee`. It's gradient-free and invariant to
+affine transforms of parameter space, so a badly-conditioned (elongated,
+correlated) target is sampled as easily as an isotropic one -- no tuning.
+
+MLX transforms on display:
+  - mx.vmap     : batch a single-point log-density over all walkers
+  - mx.compile  : fuse the stretch-move sweep into one graph
+  - mx.random   : JAX-style functional keys (split per step), so the
+                  compiled step is pure
+"""
+import mlx.core as mx
+
+
+def make_sampler(logp_single, n_dim, a=2.0):
+    """Return a compiled half-ensemble update for the G&W stretch move."""
+    logp = mx.vmap(logp_single)  # (m, D) -> (m,)
+
+    @mx.compile
+    def update_half(active, complement, key):
+        m = active.shape[0]
+        k_part, k_z, k_acc = mx.random.split(key, 3)
+        # Each active walker picks a partner from the *complementary* half.
+        j = mx.random.randint(0, complement.shape[0], (m,), key=k_part)
+        partners = mx.take(complement, j, axis=0)
+        # Stretch factor z ~ g(z) ∝ 1/sqrt(z) on [1/a, a].
+        u = mx.random.uniform(shape=(m,), key=k_z)
+        z = ((a - 1.0) * u + 1.0) ** 2 / a
+        proposal = partners + z[:, None] * (active - partners)
+        # Metropolis accept with the (D-1) stretch Jacobian.
+        log_ratio = (n_dim - 1) * mx.log(z) + logp(proposal) - logp(active)
+        accept = mx.log(mx.random.uniform(shape=(m,), key=k_acc)) < log_ratio
+        new_active = mx.where(accept[:, None], proposal, active)
+        return new_active, accept.sum()
+
+    return update_half
+
+
+def run_ensemble(logp_single, ensemble, n_steps, burn, key, a=2.0):
+    """Sample with the affine-invariant ensemble. Returns (flat samples, accept_frac)."""
+    n_walkers, n_dim = ensemble.shape
+    half = n_walkers // 2
+    update_half = make_sampler(logp_single, n_dim, a)
+
+    chain, accepted = [], mx.array(0)
+    e = ensemble
+    for t in range(n_steps):
+        key, k0, k1 = mx.random.split(key, 3)
+        h0, h1 = e[:half], e[half:]
+        h0, n0 = update_half(h0, h1, k0)      # update half 0 against half 1
+        h1, n1 = update_half(h1, h0, k1)      # update half 1 against new half 0
+        e = mx.concatenate([h0, h1], axis=0)
+        accepted = accepted + n0 + n1
+        mx.eval(e, accepted)                  # keep the lazy graph shallow
+        if t >= burn:
+            chain.append(e)
+    samples = mx.stack(chain, axis=0).reshape(-1, n_dim)
+    accept_frac = float(accepted) / (n_steps * n_walkers)
+    return samples, accept_frac

mlxmc/hmc.py

@@ -0,0 +1,59 @@
+"""Hamiltonian Monte Carlo in MLX, with an identity mass matrix.
+
+MLX transforms on display:
+  - mx.grad  : gradient of the single-point log-density (the thing HMC needs
+               and the ensemble sampler didn't)
+  - mx.vmap  : compose grad ∘ vmap to batch the gradient over all chains
+  - mx.compile : fuse the L-step leapfrog + Metropolis accept into one graph
+
+Identity mass (no preconditioning), so the contrast with the affine-invariant
+ensemble on an ill-conditioned target is visible. See `preconditioned` and
+`warmup` for the mass-matrix versions.
+"""
+import mlx.core as mx
+
+
+def make_hmc(logp_single, eps, n_leap):
+    grad_logp = mx.vmap(mx.grad(logp_single))   # (n, D) -> (n, D)
+    logp = mx.vmap(logp_single)                 # (n, D) -> (n,)
+
+    @mx.compile
+    def step(q, key):
+        n, _ = q.shape
+        k_p, k_acc = mx.random.split(key, 2)
+        p0 = mx.random.normal(shape=q.shape, key=k_p)   # resample momentum ~ N(0, I)
+        logp_q = logp(q)
+
+        # Leapfrog: half-kick, then L drifts with full-kicks between, final half-kick.
+        qq = q
+        p = p0 + 0.5 * eps * grad_logp(qq)
+        for i in range(n_leap):
+            qq = qq + eps * p
+            if i != n_leap - 1:
+                p = p + eps * grad_logp(qq)
+        p = p + 0.5 * eps * grad_logp(qq)
+
+        # Metropolis on the Hamiltonian H = -logp + 0.5 |p|^2.
+        logp_new = logp(qq)
+        log_accept = (logp_new - logp_q) + 0.5 * ((p0 * p0).sum(1) - (p * p).sum(1))
+        accept = mx.log(mx.random.uniform(shape=(n,), key=k_acc)) < log_accept
+        q_new = mx.where(accept[:, None], qq, q)
+        return q_new, accept.sum()
+
+    return step
+
+
+def run_hmc(logp_single, q0, n_steps, burn, eps, n_leap, key):
+    """Sample with fixed-step, fixed-L HMC. Returns (flat samples, accept_frac)."""
+    step = make_hmc(logp_single, eps, n_leap)
+    chain, accepted = [], mx.array(0)
+    q = q0
+    for t in range(n_steps):
+        key, k = mx.random.split(key, 2)
+        q, na = step(q, k)
+        accepted = accepted + na
+        mx.eval(q, accepted)
+        if t >= burn:
+            chain.append(q)
+    samples = mx.stack(chain, axis=0).reshape(-1, q0.shape[1])
+    return samples, float(accepted) / (n_steps * q0.shape[0])

mlxmc/nuts.py

@@ -0,0 +1,211 @@
+"""NUTS (No-U-Turn Sampler), multinomial variant (Hoffman & Gelman 2014 + Betancourt 2017),
+vectorized over chains in MLX.
+
+The MLX story (Phase 2): MLX has no `while_loop`/`scan` (see the README's "Why MLX"), so the tree-doubling
+recursion runs in **host Python** while each leapfrog leaf is `vmap`'d over all chains and
+`mx.compile`'d. Chains U-turn at different depths; a finished chain still rides along in the
+batched leapfrog but is **masked out** (`mx.where` on a per-chain `cont` flag), so it's frozen
+correctly. The doubling loop stops at `max_tree_depth` or when no chain is still going. The gap
+vs JAX's traced `while_loop` is the wasted leapfrogs on already-stopped chains (mean tree depth
+<< max is the tell). Trajectory length is adaptive, which removes the fixed-L resonance that
+eps-jitter papers over for fixed-L HMC.
+
+`eps` is a per-step argument to the returned step (not closed over) so dual-averaging in
+`nuts_warmup` can vary it every iteration without rebuilding or recompiling the leapfrog leaf.
+`Minv` is closed over for the compile, so changing M (at window boundaries during warmup)
+means rebuilding the closure -- a few recompile events over a warmup, not per step.
+"""
+import mlx.core as mx
+import numpy as np
+
+from mlxmc.warmup import DualAveraging, regularize_cov, stan_windows
+
+DMAX = 1000.0      # divergence threshold on the Hamiltonian error
+NEG = -1e30        # stand-in for log-weight 0 (divergent leaf); finite to keep logaddexp NaN-free
+
+
+def expit(x):                                  # logistic; finite NEG keeps this NaN-free
+    return 1.0 / (1.0 + mx.exp(-x))
+
+
+def logaddexp(a, b):
+    m = mx.maximum(a, b)
+    return m + mx.log(mx.exp(a - m) + mx.exp(b - m))
+
+
+def wsel(mask, a, b):                          # per-chain select; broadcasts (N,) over (N,D)
+    return mx.where(mask[:, None] if a.ndim == 2 else mask, a, b)
+
+
+def make_nuts(logp_single, Minv_np, max_tree_depth=10):
+    """Build a NUTS step `step(theta, key, eps)` that returns (sample, depths, mean_accept).
+
+    `mean_accept` is the H&G Alg 6 dual-averaging statistic -- the per-leaf Metropolis
+    acceptance min(1, exp(-Delta_H)) averaged over all leapfrog leaves in the iteration's
+    tree (per chain), then averaged over chains. Zeroed on divergent leaves. Returned for
+    use by `nuts_warmup`; `run_nuts` discards it.
+    """
+    grad_logp = mx.vmap(mx.grad(logp_single))
+    logp = mx.vmap(logp_single)
+    Minv = mx.array(Minv_np.astype(np.float32))
+    Mhalf_T = mx.array(np.linalg.cholesky(np.linalg.inv(Minv_np)).T.astype(np.float32))
+
+    @mx.compile
+    def leap(theta, r, se):                    # one leapfrog; se = signed step, (N,1)
+        r = r + 0.5 * se * grad_logp(theta)
+        theta = theta + se * (r @ Minv)
+        r = r + 0.5 * se * grad_logp(theta)
+        return theta, r
+
+    def joint(theta, r):                        # -H = logp - 0.5 r^T M^-1 r  (the log-weight)
+        return logp(theta) - 0.5 * ((r @ Minv) * r).sum(1)
+
+    def no_uturn(tm, rm, tp, rp):               # True = keep going (no U-turn), generalized metric
+        d = tp - tm
+        return ((d * (rm @ Minv)).sum(1) >= 0) & ((d * (rp @ Minv)).sum(1) >= 0)
+
+    def build(theta, r, lw0, depth, dirn, eps, key):
+        """Recursive subtree builder. Returns the two endpoints, the multinomial proposal from
+        this subtree, its total log-weight, a per-chain continue flag, and the leaf-accept
+        (sum, count) used by the dual-averaging statistic."""
+        if depth == 0:                          # base: a single leapfrog in `dirn`
+            th1, r1 = leap(theta, r, (dirn * eps)[:, None])
+            lw1 = joint(th1, r1)
+            div = (lw0 - lw1 > DMAX) | mx.isnan(lw1)
+            # H&G Alg 6 leaf statistic: min(1, exp(lw1 - lw0)) = min(1, exp(-Delta_H)); zeroed
+            # on divergence/NaN. exp(min(0, ...)) keeps the value <= 1 without an explicit clamp.
+            leaf_a = mx.where(div, mx.zeros_like(lw1),
+                              mx.exp(mx.minimum(mx.zeros_like(lw1), lw1 - lw0)))
+            ones = mx.ones_like(lw1)
+            return th1, r1, th1, r1, th1, mx.where(div, NEG, lw1), ~div, leaf_a, ones
+
+        kL, kR, ks = mx.random.split(key, 3)
+        tm, rm, tp, rp, p1, lw1, s1, a1, c1 = build(theta, r, lw0, depth - 1, dirn, eps, kL)
+        lt, lr = wsel(dirn < 0, tm, tp), wsel(dirn < 0, rm, rp)      # extend the leading edge
+        tm2, rm2, tp2, rp2, p2, lw2, s2, a2, c2 = build(lt, lr, lw0, depth - 1, dirn, eps, kR)
+
+        ftm, frm = wsel(dirn < 0, tm2, tm), wsel(dirn < 0, rm2, rm)  # stitched full endpoints
+        ftp, frp = wsel(dirn < 0, tp, tp2), wsel(dirn < 0, rp, rp2)
+        pick2 = mx.random.uniform(shape=(theta.shape[0],), key=ks) < expit(lw2 - lw1)
+        prop = wsel(pick2, p2, p1)
+        lw = logaddexp(lw1, lw2)
+        s = s1 & s2 & no_uturn(ftm, frm, ftp, frp)
+        # Accept stat sums all leaves traversed (H&G: n_alpha counts every leapfrog leaf,
+        # not just the ones in valid subtrees). State propagation still gates on s1 as before.
+        return (wsel(s1, ftm, tm), wsel(s1, frm, rm), wsel(s1, ftp, tp), wsel(s1, frp, rp),
+                wsel(s1, prop, p1), mx.where(s1, lw, lw1), s1 & s, a1 + a2, c1 + c2)
+
+    def step(theta, key, eps):
+        N = theta.shape[0]
+        km, k = mx.random.split(key, 2)
+        r0 = mx.random.normal(shape=theta.shape, key=km) @ Mhalf_T   # ~ N(0, M)
+        lw0 = joint(theta, r0)
+        tm = tp = theta
+        rm = rp = r0
+        sample, lw_tree = theta, lw0
+        cont = mx.array(np.ones(N, dtype=bool))
+        depths = mx.zeros((N,), dtype=mx.int32)
+        accept_sum = mx.zeros((N,))
+        accept_cnt = mx.zeros((N,))
+
+        for depth in range(max_tree_depth):
+            cont_was = cont                     # gate accept-stat accumulation on entering state
+            k, kdir, ksub, ksel = mx.random.split(k, 4)
+            dirn = mx.where(mx.random.uniform(shape=(N,), key=kdir) < 0.5, -1.0, 1.0)
+            depths = depths + cont.astype(mx.int32)
+            lt, lr = wsel(dirn < 0, tm, tp), wsel(dirn < 0, rm, rp)
+            ntm, nrm, ntp, nrp, prop, lw_sub, s_sub, a_sub, c_sub = build(
+                lt, lr, lw0, depth, dirn, eps, ksub)
+
+            new_tm, new_rm = wsel(dirn < 0, ntm, tm), wsel(dirn < 0, nrm, rm)
+            new_tp, new_rp = wsel(dirn < 0, tp, ntp), wsel(dirn < 0, rp, nrp)
+            # multinomial: adopt the new subtree's proposal with prob W_sub / (W_tree + W_sub),
+            # but ONLY if the subtree is valid (H&G Alg 3: gate on s'). Adopting proposals from
+            # a subtree that internally U-turned/diverged over-samples its returned far points.
+            pick = (mx.random.uniform(shape=(N,), key=ksel) < expit(lw_sub - lw_tree)) & cont & s_sub
+            sample = wsel(pick, prop, sample)
+            lw_tree = mx.where(cont, logaddexp(lw_tree, lw_sub), lw_tree)
+            tm, rm = wsel(cont, new_tm, tm), wsel(cont, new_rm, rm)
+            tp, rp = wsel(cont, new_tp, tp), wsel(cont, new_rp, rp)
+            # Chains that had already stopped contribute no new leaves this iteration.
+            accept_sum = accept_sum + mx.where(cont_was, a_sub, mx.zeros_like(a_sub))
+            accept_cnt = accept_cnt + mx.where(cont_was, c_sub, mx.zeros_like(c_sub))
+            cont = cont & s_sub & no_uturn(new_tm, new_rm, new_tp, new_rp)
+            mx.eval(cont, sample, tm, tp, rm, rp, lw_tree, depths, accept_sum, accept_cnt)
+            if cont.sum().item() == 0:          # host-side early stop once all chains U-turned
+                break
+
+        # Per-chain leaf-mean accept, then mean over chains. Empty count -> 1 to avoid 0/0.
+        per_chain = accept_sum / mx.maximum(accept_cnt, mx.ones_like(accept_cnt))
+        return sample, depths, per_chain.mean()
+
+    return step
+
+
+def run_nuts(logp_single, theta0, n_samples, eps, Minv_np, key, max_tree_depth=10):
+    """Returns (chain (T,N,D), mean_tree_depth, max_tree_depth_seen). theta0 already warmed.
+    max-vs-mean depth is the masking-overhead tell: the batch pays the deepest chain's cost."""
+    step = make_nuts(logp_single, Minv_np, max_tree_depth)
+    chain, depth_sum, depth_max, theta = [], 0.0, 0, theta0
+    for _ in range(n_samples):
+        key, k = mx.random.split(key, 2)
+        theta, depths, _ = step(theta, k, eps)    # discard the leaf-accept stat for sampling
+        mx.eval(theta, depths)
+        chain.append(theta)
+        depth_sum += float(depths.mean())
+        depth_max = max(depth_max, int(depths.max()))
+    return mx.stack(chain, axis=0), depth_sum / n_samples, depth_max
+
+
+def nuts_warmup(logp_single, q0, n_warmup, key, eps0=0.25, target_accept=0.8,
+                init_buffer=75, term_buffer=50, base_window=25, max_tree_depth=10):
+    """NUTS-specific warmup: dual-averaging on NUTS's tree-averaged leaf-acceptance statistic,
+    Stan-style windowed dense-M estimation. Returns (q_last, eps_bar, Minv_np) ready for
+    `run_nuts` -- the same interface as `mlxmc.warmup.warmup`, but the tuned eps reflects
+    NUTS's adaptive trajectory length rather than borrowing fixed-L HMC's optimum.
+
+    The leaf-accept stat (H&G 2014 Alg 6): per leapfrog leaf, alpha_leaf = min(1, exp(-Delta_H)),
+    zeroed on divergence; alpha for the iteration is the mean over leaves per chain, then over
+    chains. Stan defaults target this at 0.8 (the same number used here -- the statistic differs,
+    so the resulting eps does too).
+
+    Implementation notes:
+      - `eps` is a per-step argument to the NUTS step, so dual-averaging changes it every
+        iteration without recompiling the leapfrog leaf.
+      - `Minv` is closed over the compiled leaf, so when the windowed estimator produces a
+        new M, the NUTS step is rebuilt (one recompile per window boundary, not per step).
+      - Covariance + Cholesky are done host-side in fp64; only the leapfrog runs fp32.
+      - A NaN accept stat (e.g. all leaves diverged in a single iteration) is treated as 0
+        so dual-averaging stays finite.
+    """
+    n_chains, d = q0.shape
+    init_buffer, term_buffer, window_ends = stan_windows(
+        n_warmup, init_buffer, term_buffer, base_window)
+    boundaries = set(window_ends)
+
+    da = DualAveraging(eps0, target_accept)
+    Minv_np = np.eye(d)
+    step = make_nuts(logp_single, Minv_np, max_tree_depth)
+    eps = eps0
+
+    q, window_samples = q0, []
+    for t in range(n_warmup):
+        key, k = mx.random.split(key, 2)
+        q, _depths, accept_prob = step(q, k, eps)
+        mx.eval(q, accept_prob)
+        a = float(accept_prob)
+        if not np.isfinite(a):
+            a = 0.0                                       # treat all-divergent batch as reject
+        eps = da.update(a)
+
+        if init_buffer <= t < (n_warmup - term_buffer):
+            window_samples.append(np.array(q))            # collect cov samples in slow windows
+        if (t + 1) in boundaries and window_samples:
+            X = np.concatenate(window_samples, axis=0)    # (steps * n_chains, d)
+            Minv_np = regularize_cov(np.cov(X, rowvar=False), X.shape[0])
+            step = make_nuts(logp_single, Minv_np, max_tree_depth)   # rebuild for the new M
+            window_samples = []
+            da.restart(da.eps_bar)                        # re-anchor eps to the new metric
+            eps = da.eps_bar
+
+    return q, da.eps_bar, Minv_np

mlxmc/preconditioned.py

@@ -0,0 +1,56 @@
+"""Preconditioned HMC: mass matrix M = Sigma^{-1} makes the dynamics isotropic.
+
+This closes the loop from the affine-invariance discussion: the mass matrix is
+HMC's version of affine invariance -- but you must *supply* it (here you'd pass
+the true Sigma; in practice you'd estimate it during warmup, as NUTS/Stan do).
+With the right M, HMC mixes with far fewer, cheaper leapfrog steps.
+"""
+import mlx.core as mx
+
+
+def make_phmc(logp_single, eps, n_leap, Minv, Mhalf):
+    grad_logp = mx.vmap(mx.grad(logp_single))
+    logp = mx.vmap(logp_single)
+    Minv = mx.array(Minv)          # M^{-1} = Sigma   (drift + kinetic energy)
+    Mhalf_T = mx.transpose(mx.array(Mhalf))   # chol(Sigma^{-1})^T  (momentum draw)
+
+    def kinetic(p):                # 0.5 p^T M^{-1} p
+        return 0.5 * ((p @ Minv) * p).sum(1)
+
+    @mx.compile
+    def step(q, key):
+        n, _ = q.shape
+        k_p, k_acc = mx.random.split(key, 2)
+        z = mx.random.normal(shape=q.shape, key=k_p)
+        p0 = z @ Mhalf_T           # ~ N(0, M)
+        logp_q, K0 = logp(q), kinetic(p0)
+
+        qq = q
+        p = p0 + 0.5 * eps * grad_logp(qq)
+        for i in range(n_leap):
+            qq = qq + eps * (p @ Minv)          # drift uses M^{-1} = Sigma
+            if i != n_leap - 1:
+                p = p + eps * grad_logp(qq)
+        p = p + 0.5 * eps * grad_logp(qq)
+
+        log_accept = (logp(qq) - logp_q) + (K0 - kinetic(p))
+        accept = mx.log(mx.random.uniform(shape=(n,), key=k_acc)) < log_accept
+        return mx.where(accept[:, None], qq, q), accept.sum()
+
+    return step
+
+
+def run_phmc(logp_single, q0, n_steps, burn, eps, n_leap, key, Minv, Mhalf):
+    """Sample with preconditioned (mass-matrix) HMC. Returns the structured (T, N, D) chain.
+
+    `Minv` is M^{-1} (= the covariance you precondition with) and `Mhalf` is chol(M).
+    """
+    step = make_phmc(logp_single, eps, n_leap, Minv, Mhalf)
+    chain, q = [], q0
+    for t in range(n_steps):
+        key, k = mx.random.split(key, 2)
+        q, _ = step(q, k)
+        mx.eval(q)
+        if t >= burn:
+            chain.append(q)
+    return mx.stack(chain, axis=0)

mlxmc/targets.py

@@ -0,0 +1,70 @@
+"""Example / benchmark target distributions, each a single-point log-density
+`logp(x) -> scalar` suitable for any sampler in mlxmc.
+
+  - Gaussian: strongly correlated, ill-conditioned -- easy for the affine-invariant
+    ensemble, hard for identity-mass HMC. The canonical correctness target.
+  - Banana (Haario twisted Gaussian): a curved ridge; local covariance rotates along it,
+    so no single linear preconditioner (constant M) or affine map is right everywhere.
+  - Funnel (Neal): v ~ N(0, 3^2), x_i | v ~ N(0, exp(v)). The scale of x spans orders of
+    magnitude with v, so a constant M is wrong everywhere at once and the neck (v << 0)
+    has stiff gradients that diverge in fp32. The honest fix is geometry-aware coordinates.
+  - Non-centered funnel: sample (v, x̃) with x = x̃·exp(v/2); the v-dependent scale drops
+    out, leaving a product of independent Gaussians -- the reparametrization that fixes HMC.
+
+The TRUTH dicts give (dim index, true mean, true std) for the cleanest diagnostics.
+"""
+import mlx.core as mx
+import numpy as np
+
+# --------------------------------------------------------------- correlated Gaussian
+GAUSSIAN_MU = np.array([1.0, -2.0])
+GAUSSIAN_SIGMA = np.array([[25.0, 4.5], [4.5, 1.0]])   # corr 0.9, 25:1 variance ratio
+_gauss_mu = mx.array(GAUSSIAN_MU)
+_gauss_sig_inv = mx.array(np.linalg.inv(GAUSSIAN_SIGMA))
+
+
+def gaussian_logp(x):                          # x: (D,) -> scalar
+    d = x - _gauss_mu
+    return -0.5 * (d @ _gauss_sig_inv @ d)
+
+
+# --------------------------------------------------------------- banana (Haario)
+B_BANANA = 0.05   # curvature; Var[x2] = 1 + 2 B^2 * 100^2
+
+
+def banana_logp(x):
+    """phi(x) = (x1, x2 + B x1^2 - 100 B, x3, ...) ~ N(0, diag(100, 1, 1, ...))."""
+    x1, x2 = x[0], x[1]
+    twisted = x2 + B_BANANA * x1 * x1 - 100.0 * B_BANANA
+    rest = 0.5 * (x[2:] * x[2:]).sum() if x.shape[0] > 2 else 0.0
+    return -(x1 * x1) / 200.0 - 0.5 * twisted * twisted - rest
+
+
+BANANA_TRUTH = {
+    "x1": (0, 0.0, 10.0),                                  # N(0, 100)
+    "x2": (1, 0.0, np.sqrt(1.0 + 2.0 * B_BANANA**2 * 100.0**2)),
+}
+
+
+# --------------------------------------------------------------- funnel (Neal)
+def funnel_logp(z):
+    """v = z[0] ~ N(0, 9); x = z[1:] | v ~ N(0, exp(v)). The -0.5*(D-1)*v term is the
+    v-dependent normalization -- it's what makes this a funnel, not a free-floating v."""
+    v, x = z[0], z[1:]
+    n_x = z.shape[0] - 1
+    return -(v * v) / 18.0 - 0.5 * n_x * v - 0.5 * mx.exp(-v) * (x * x).sum()
+
+
+def funnel_nc_logp(z):
+    """Non-centered funnel: sample (v, x̃) with x̃ ~ N(0,1), and x = x̃·exp(v/2). In these
+    coordinates the v-dependent scale vanishes from the density, leaving a *product of
+    independent Gaussians* (v ~ N(0,9), x̃ ~ N(0,1)) -- no funnel geometry, so HMC's global
+    metric is now correct everywhere. v's marginal is unchanged (N(0,9)), so the same
+    truth/diagnostics apply, and this should flip the centered-funnel result."""
+    v, xt = z[0], z[1:]
+    return -(v * v) / 18.0 - 0.5 * (xt * xt).sum()
+
+
+FUNNEL_TRUTH = {
+    "v": (0, 0.0, 3.0),                                    # N(0, 9): the honest mixing test
+}

mlxmc/warmup.py

@@ -0,0 +1,178 @@
+"""Warmup adaptation for HMC: dual-averaging step size + windowed dense mass-matrix
+estimation (Stan-style). This makes the preconditioned-HMC result *honest* -- rather
+than being handed the true Sigma, we estimate M^{-1} = Cov(q) during warmup and adapt
+eps to a target acceptance, then sample. The mass matrix is HMC's affine invariance, and
+this is how you earn it instead of supplying it.
+
+MLX notes:
+  - eps, M^{-1}, and chol(M)^T are passed as *array arguments* to the compiled step, not
+    closed-over constants (as make_hmc/make_phmc do). mx.compile recompiles on shape/dtype
+    or baked-constant changes, NOT on array-value changes -- so we can vary eps/M every
+    warmup iteration and reuse the one compiled graph.
+  - n_leap stays a Python int: it's the unrolled leapfrog length, structural to the graph.
+    Phase 1 adapts eps + M with L fixed; dynamic trajectory length is NUTS (Phase 2), which
+    this MLX has no control-flow primitive (while_loop/scan) for.
+  - Covariance + Cholesky run host-side in numpy fp64, so the ill-conditioned linear algebra
+    (the real fp32 limit) never touches the GPU; only the leapfrog runs fp32.
+  - Covariance is pooled across chains AND steps within a window (many chains => fast
+    estimate). Early overdispersion is discarded by the windowed schedule; the final, longest
+    slow window estimates M near stationarity.
+"""
+import math
+
+import mlx.core as mx
+import numpy as np
+
+
+class DualAveraging:
+    """Nesterov dual averaging for step-size adaptation (Hoffman & Gelman 2014, Alg. 5).
+
+    The chain runs on the raw `exp(log_eps)` (keeps exploring); the *averaged* `eps_bar`
+    is what we freeze for sampling. `restart` re-anchors at a window boundary, since the
+    stable step size changes when the mass matrix does.
+    """
+
+    def __init__(self, eps0, target_accept=0.8, gamma=0.05, t0=10.0, kappa=0.75):
+        self.target, self.gamma, self.t0, self.kappa = target_accept, gamma, t0, kappa
+        self.restart(eps0)
+
+    def restart(self, eps0):
+        self.mu = math.log(10.0 * eps0)          # shrink toward 10x the anchor
+        self.Hbar = 0.0
+        self.log_eps = math.log(eps0)
+        self.log_eps_bar = math.log(eps0)
+        self.m = 0
+
+    def update(self, accept_prob):
+        self.m += 1
+        m = self.m
+        w = 1.0 / (m + self.t0)
+        self.Hbar = (1.0 - w) * self.Hbar + w * (self.target - accept_prob)
+        self.log_eps = self.mu - math.sqrt(m) / self.gamma * self.Hbar
+        eta = m ** (-self.kappa)
+        self.log_eps_bar = eta * self.log_eps + (1.0 - eta) * self.log_eps_bar
+        return math.exp(self.log_eps)
+
+    @property
+    def eps_bar(self):
+        return math.exp(self.log_eps_bar)
+
+
+def stan_windows(n_warmup, init_buffer=75, term_buffer=50, base_window=25):
+    """Stan-style warmup schedule. Returns (init_buffer, term_buffer, window_ends), where
+    window_ends are iteration counts (1-based, inclusive) at which to re-estimate M and
+    restart step-size adaptation. M is held fixed during the init buffer (find the mode)
+    and term buffer (final eps polish); the slow windows between double in length, and the
+    last absorbs any remainder. Not bit-identical to Stan, but the same structure."""
+    if init_buffer + term_buffer + base_window > n_warmup:
+        # Too short for the default buffers: fall back to Stan's 15/75/10 proportions.
+        init_buffer = max(1, int(round(0.15 * n_warmup)))
+        term_buffer = max(1, int(round(0.10 * n_warmup)))
+        base_window = max(1, n_warmup - init_buffer - term_buffer)
+    last = n_warmup - term_buffer
+    ends, start, window = [], init_buffer, base_window
+    while start + window < last:
+        nxt = start + window
+        if nxt + 2 * window > last:     # next doubled window would overrun -> absorb now
+            nxt = last
+        ends.append(nxt)
+        start, window = nxt, window * 2
+    if not ends or ends[-1] != last:
+        ends.append(last)
+    return init_buffer, term_buffer, ends
+
+
+def regularize_cov(cov, n):
+    """Stan's dense-metric shrinkage toward a small diagonal (stabilizes small-n windows)."""
+    d = cov.shape[0]
+    return (n / (n + 5.0)) * cov + 1e-3 * (5.0 / (n + 5.0)) * np.eye(d)
+
+
+def make_warmup_step(logp_single, n_leap):
+    """Leapfrog + Metropolis step with eps, M^{-1}, chol(M)^T as array arguments.
+    Returns (q_new, mean_accept_prob); the continuous accept prob feeds dual averaging."""
+    grad_logp = mx.vmap(mx.grad(logp_single))
+    logp = mx.vmap(logp_single)
+
+    @mx.compile
+    def step(q, key, eps, Minv, Mhalf_T):
+        n, _ = q.shape
+        k_p, k_acc = mx.random.split(key, 2)
+        z = mx.random.normal(shape=q.shape, key=k_p)
+        p0 = z @ Mhalf_T                          # ~ N(0, M)
+        K0 = 0.5 * ((p0 @ Minv) * p0).sum(1)      # 0.5 p^T M^{-1} p
+        logp_q = logp(q)
+
+        qq = q
+        p = p0 + 0.5 * eps * grad_logp(qq)
+        for i in range(n_leap):
+            qq = qq + eps * (p @ Minv)            # drift uses M^{-1}
+            if i != n_leap - 1:
+                p = p + eps * grad_logp(qq)
+        p = p + 0.5 * eps * grad_logp(qq)
+
+        K = 0.5 * ((p @ Minv) * p).sum(1)
+        log_accept = (logp(qq) - logp_q) + (K0 - K)
+        accept_prob = mx.minimum(1.0, mx.exp(log_accept))     # continuous, for dual avg
+        # A divergent leapfrog (NaN energy, e.g. the funnel neck in fp32) must REJECT, not
+        # poison the dual-averaging mean with NaN. Treat NaN accept prob as 0.
+        accept_prob = mx.where(mx.isnan(accept_prob), mx.zeros_like(accept_prob), accept_prob)
+        accept = mx.random.uniform(shape=(n,), key=k_acc) < accept_prob
+        q_new = mx.where(accept[:, None], qq, q)
+        return q_new, accept_prob.mean()
+
+    return step
+
+
+def warmup(logp_single, q0, n_warmup, n_leap, key, eps0=0.25, target_accept=0.8,
+           init_buffer=75, term_buffer=50, base_window=25):
+    """Run Stan-style warmup. Returns (q_last, eps_bar, Minv_np) -- the tuned step size and
+    estimated M^{-1} = Cov(q), ready to hand to run_chain."""
+    n_chains, d = q0.shape
+    init_buffer, term_buffer, window_ends = stan_windows(
+        n_warmup, init_buffer, term_buffer, base_window)
+    boundaries = set(window_ends)
+
+    da = DualAveraging(eps0, target_accept)
+    step = make_warmup_step(logp_single, n_leap)
+
+    Minv_np = np.eye(d)                                   # start from the identity metric
+    eps = mx.array(eps0, dtype=mx.float32)
+    Minv = mx.array(Minv_np.astype(np.float32))
+    Mhalf_T = mx.array(np.linalg.cholesky(np.linalg.inv(Minv_np)).T.astype(np.float32))
+
+    q, window_samples = q0, []
+    for t in range(n_warmup):
+        key, k = mx.random.split(key, 2)
+        q, acc_prob = step(q, k, eps, Minv, Mhalf_T)
+        mx.eval(q, acc_prob)
+        eps = mx.array(da.update(float(acc_prob)), dtype=mx.float32)   # adapt eps every step
+
+        if init_buffer <= t < (n_warmup - term_buffer):
+            window_samples.append(np.array(q))            # collect cov samples in slow windows
+        if (t + 1) in boundaries and window_samples:
+            X = np.concatenate(window_samples, axis=0)    # (steps * n_chains, d)
+            Minv_np = regularize_cov(np.cov(X, rowvar=False), X.shape[0])
+            Minv = mx.array(Minv_np.astype(np.float32))
+            Mhalf_T = mx.array(np.linalg.cholesky(np.linalg.inv(Minv_np)).T.astype(np.float32))
+            window_samples = []
+            da.restart(da.eps_bar)                        # re-anchor eps to the new metric
+            eps = mx.array(da.eps_bar, dtype=mx.float32)
+
+    return q, da.eps_bar, Minv_np
+
+
+def run_chain(logp_single, q0, n_steps, burn, eps, Minv_np, key, n_leap):
+    """Sample with fixed tuned (eps, M). Returns the structured (T, N, D) chain for ESS."""
+    step = make_warmup_step(logp_single, n_leap)
+    eps_a = mx.array(eps, dtype=mx.float32)
+    Minv_a = mx.array(Minv_np.astype(np.float32))
+    Mhalf_T = mx.array(np.linalg.cholesky(np.linalg.inv(Minv_np)).T.astype(np.float32))
+    chain, q = [], q0
+    for t in range(n_steps):
+        key, k = mx.random.split(key, 2)
+        q, _ = step(q, k, eps_a, Minv_a, Mhalf_T)
+        mx.eval(q)
+        if t >= burn:
+            chain.append(q)
+    return mx.stack(chain, axis=0)

mlxmc-0.1.0.dist-info/METADATA

@@ -0,0 +1,200 @@
+Metadata-Version: 2.4
+Name: mlxmc
+Version: 0.1.0
+Summary: MCMC samplers in Apple MLX
+Project-URL: Homepage, https://github.com/jrcheshire/mlxmc
+Project-URL: Repository, https://github.com/jrcheshire/mlxmc
+Project-URL: Issues, https://github.com/jrcheshire/mlxmc/issues
+Author-email: Jamie Cheshire <cheshire@caltech.edu>
+License-Expression: BSD-3-Clause
+License-File: LICENSE
+Keywords: apple-silicon,bayesian,ensemble-sampler,hmc,mcmc,mlx,nuts,sampling
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: MacOS
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.11
+Requires-Dist: mlx<0.30,>=0.29.3
+Requires-Dist: numpy<3,>=2
+Provides-Extra: viz
+Requires-Dist: matplotlib<4,>=3.10; extra == 'viz'
+Description-Content-Type: text/markdown
+
+# mlxmc
+
+MCMC samplers written in Apple [MLX](https://github.com/ml-explore/mlx), using its
+`grad` / `vmap` / `compile` transforms. MLX has no probabilistic-programming library
+yet (nothing like BlackJAX or NumPyro), so this is a first pass at one.
+
+> **Status: research code.** The samplers are tested (moment recovery, Σ-estimation,
+> affine invariance, and the autocorrelation diagnostics, on both the CPU and Metal
+> backends), but the API is young and likely to change.
+
+## What's here
+
+The package lives under `src/mlxmc/`; runnable demos and the benchmark study are in
+`examples/`.
+
+| Module (`mlxmc.`) | Sampler / tool |
+|---|---|
+| `ensemble` | Affine-invariant ensemble (Goodman & Weare 2010 — the `emcee` stretch move). Gradient-free, tuning-free. `make_sampler`, `run_ensemble`. |
+| `hmc` | Hamiltonian Monte Carlo, identity mass. `grad ∘ vmap` batched over chains. `make_hmc`, `run_hmc`. |
+| `preconditioned` | Mass-matrix HMC (M = Σ⁻¹). `make_phmc`, `run_phmc`. |
+| `warmup` | Stan-style warmup: dual-averaging step size + windowed **dense** mass-matrix estimation. `warmup`, `run_chain`. |
+| `nuts` | NUTS (multinomial; Hoffman & Gelman 2014), vectorized over chains. `make_nuts`, `run_nuts`. |
+| `diagnostics` | Effective sample size / integrated autocorrelation time (FFT + Sokal window); the cross-sampler **ESS/sec** metric. |
+| `targets` | Example log-densities: correlated Gaussian, banana, centered / non-centered funnel, with known moments. |
+
+| Example (`examples/`) | What it shows |
+|---|---|
+| `gaussian_ess.py` | Ensemble vs identity-mass HMC vs preconditioned HMC by ESS/sec on the Gaussian. |
+| `warmup_validation.py` | Warmup recovers the true Σ and matches oracle ESS/sec. |
+| `hard_targets.py` | Banana + funnel benchmark (`lscan` / `dscan` modes). |
+| `nuts_funnel.py` | NUTS correctness on the Gaussian; `funnel` mode for the masking-overhead study. |
+| `affine_invariance.py` | Empirical proof of affine invariance (same RNG → bit-identical acceptance under an affine map). |
+| `plot_hard_targets.py` | Renders `hard_targets_figure.png` (needs the optional `viz` env). |
+
+## Why MLX
+
+`grad`, `vmap`, `jvp`/`vjp`, and `compile` transfer almost directly from JAX,
+with JAX-style functional RNG keys (`mx.random.split`). The wrinkles that shape
+this code:
+
+- **No traced control-flow primitives** (no `while_loop` / `scan` / `cond`). MLX
+  is eager execution plus `compile` of *static* graphs. Fixed-length unrolled
+  loops (leapfrog, fixed-`L` HMC) compile fine; data-dependent trajectory length
+  (NUTS) is the hard case — `mlxmc.nuts` runs every chain to a fixed `max_tree_depth`
+  and **masks** finished chains.
+- **fp32 on the GPU.** Apple Metal has no fp64 in hardware (MLX has fp64 only on
+  the CPU backend). This is fine for sampling — Monte Carlo error (~1/√ESS) swamps
+  fp32 roundoff (~1e-6) — but ill-conditioned linear algebra (covariance, Cholesky
+  in warmup) is kept host-side in numpy fp64; only the leapfrog runs on the GPU.
+
+## Install
+
+This is a [pixi](https://pixi.sh) project (installs the package editable):
+
+```bash
+pixi install
+pixi run python examples/gaussian_ess.py             # ensemble vs HMC vs preconditioned
+pixi run python examples/nuts_funnel.py funnel        # several examples have demo modes
+pixi run -e viz python examples/plot_hard_targets.py  # plotting needs the optional viz env
+```
+
+Or install into any environment with pip: `pip install -e .` (needs `mlx`, so arm64
+macOS). Add the plotting extra with `pip install -e ".[viz]"` (matplotlib).
+
+## Usage
+
+Every sampler takes a single-point log-density `logp(x) -> scalar` for `x` of
+shape `(D,)`; batching over walkers/chains is handled internally with `vmap`.
+Positions are MLX arrays of shape `(n_chains, D)`.
+
+```python
+import mlx.core as mx
+import numpy as np
+
+# Target: a strongly correlated 2-D Gaussian (corr 0.9, 25:1 variance ratio).
+# mlxmc.targets ships this one (as `gaussian_logp`) plus banana / funnel.
+mu = mx.array([1.0, -2.0])
+Sig_inv = mx.array(np.linalg.inv([[25.0, 4.5], [4.5, 1.0]]))
+
+def logp(x):                              # x: (D,) -> scalar
+    d = x - mu
+    return -0.5 * (d @ Sig_inv @ d)
+
+key = mx.random.key(0)
+```
+
+**Gradient-free ensemble** — no tuning, handles the ill-conditioning for free:
+
+```python
+from mlxmc import run_ensemble
+
+key, k = mx.random.split(key)
+ensemble = mx.random.normal(shape=(2000, 2), key=k) * 5.0     # (n_walkers, D)
+samples, accept_frac = run_ensemble(logp, ensemble, n_steps=3000, burn=1000, key=key)
+```
+
+**HMC, hand-tuned**, and **NUTS after Stan-style warmup** (same `logp`):
+
+```python
+from mlxmc import run_hmc, warmup, run_nuts
+
+key, k = mx.random.split(key)
+q0 = mx.random.normal(shape=(1000, 2), key=k) * 5.0           # (n_chains, D)
+
+samples, acc = run_hmc(logp, q0, n_steps=1500, burn=500,
+                       eps=0.15, n_leap=40, key=key)
+
+# Warmup adapts (eps, dense M); NUTS then adapts trajectory length itself.
+q_last, eps, Minv = warmup(logp, q0, n_warmup=600, n_leap=8, key=key)
+chain, mean_depth, max_depth = run_nuts(logp, q_last, n_samples=1500,
+                                        eps=eps, Minv_np=Minv, key=key)
+```
+
+> **Return shapes differ by sampler.** `run_ensemble` and `run_hmc` return
+> `(samples, accept_frac)` with `samples` flattened to `(n_draws, D)`.
+> `run_phmc`, `run_chain` (post-warmup HMC), and `run_nuts` return a structured
+> `(steps, chains, D)` chain — the layout `mlxmc.diagnostics` expects for ESS —
+> and `run_nuts` additionally returns the mean/max tree depth.
+
+## Findings
+
+![Sampler benchmarks on the banana and funnel targets](https://raw.githubusercontent.com/jrcheshire/mlxmc/main/hard_targets_figure.png)
+
+Validated on a corr-0.9, 25:1-variance Gaussian and on banana / funnel targets;
+every number below is reproducible with the scripts in
+[`examples/`](https://github.com/jrcheshire/mlxmc/tree/main/examples):
+
+- **Affine-invariant ensemble** is the robust low-D default: gradient-free,
+  tuning-free, handles ill-conditioning for free (acceptance is bit-identical
+  under an affine map). But weaker per-step mixing and it degrades with dimension.
+- **HMC** needs gradients and a tuned `eps`/`L`, but mixes far better
+  (τ≈2 vs ≈26). A **warmup-adapted dense mass matrix** recovers the true Σ to
+  <1% Frobenius error and buys ~7–11× the ESS/sec — HMC's version of affine
+  invariance, earned rather than supplied.
+- **Fixed-`L` HMC has a trajectory resonance:** on near-Gaussian targets, when
+  `eps·L` lands near a multiple of 2π the trajectory returns to its start and
+  mixing collapses. Jittering `eps` per trajectory cures it; NUTS's adaptive
+  trajectory length is the principled fix.
+- **NUTS** is validated exact on the Gaussian (recovered covariance 24.97 vs 25)
+  and auto-tunes trajectory length, but vectorized NUTS pays a real masking cost
+  when trajectory lengths are heterogeneous — with no `while_loop`, every chain
+  runs to the deepest chain's tree depth, up to a ~30× wall-time penalty at the
+  funnel mouth versus the same target reparametrized.
+- **Geometry matters more than the sampler:** on the *centered* funnel the
+  gradient-free ensemble beats a global-metric HMC, because a constant mass matrix
+  is wrong everywhere when the scale is position-dependent; a **non-centered
+  reparametrization** removes the geometry and makes HMC unbiased again.
+- **ESS/sec is the honest efficiency metric** — acceptance fraction is a
+  misleading proxy.
+
+## Development
+
+```bash
+pixi run test                                   # full suite on the default device
+MLXMC_TEST_DEVICE=cpu pixi run test             # force the CPU backend
+MLXMC_TEST_DEVICE=gpu pixi run test             # force the Metal GPU
+```
+
+The suite (`tests/`) checks moment recovery for every sampler, warmup's Σ
+estimate, the affine-invariance identity, and the autocorrelation-time
+diagnostics. A GitHub Actions workflow (`.github/workflows/tests.yml`) runs the
+CPU + GPU matrix on an Apple-silicon runner for pull requests to `main` (and on
+manual dispatch from the Actions tab). Direct pushes to `main` don't trigger it,
+which keeps the (10x-billed) macOS runner minutes down.
+
+## References
+
+- Goodman & Weare (2010), *Ensemble samplers with affine invariance.*
+- Hoffman & Gelman (2014), *The No-U-Turn Sampler.*
+- Betancourt (2017), *A Conceptual Introduction to Hamiltonian Monte Carlo.*
+
+## License
+
+[BSD-3-Clause](https://github.com/jrcheshire/mlxmc/blob/main/LICENSE).

mlxmc-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+mlxmc/__init__.py,sha256=hRv4Sm6NAklNuG43DnCsaz5KhRh-zRxP_oZ7XDp1FxI,1026
+mlxmc/diagnostics.py,sha256=vET8WYBdP5HK_pr2WGjbrMZbPbA5XV9iCcj7RzzQR2E,1816
+mlxmc/ensemble.py,sha256=t04msVDGj3B_-Z8_Ifyw8fVJ_MWlnoPrpfBqtr2dVfU,2617
+mlxmc/hmc.py,sha256=Vjmop3-X7FeAVqZVnKjPC7PG6UwXk_puhG5W-wju0X0,2242
+mlxmc/nuts.py,sha256=I-TYCDl-vj8BNr0sDZfGxBu1IUGvjRAShLcLJEtGvKY,11342
+mlxmc/preconditioned.py,sha256=hTQkmmNgMW5rgQ4ajlok1dQLc8a5dcBqg5Be-9oehow,2113
+mlxmc/targets.py,sha256=JD95uJJ3oZzYo1vAGB20aHCtGEBR1TYZGrID215FiUo,3165
+mlxmc/warmup.py,sha256=MPAM9otpmNT8PelOARLr_bMI3qbKwhaEnFKnf7uLPu0,8235
+mlxmc-0.1.0.dist-info/METADATA,sha256=hwytPlj-kBRqcEhQ10Wlmo4Mg730XlazKdPPrXaYgns,9499
+mlxmc-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+mlxmc-0.1.0.dist-info/licenses/LICENSE,sha256=k4YNKXQbOe0QvMkczCNPA_bck_92SZ5gNsnRbpA23ko,1501
+mlxmc-0.1.0.dist-info/RECORD,,

mlxmc-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

mlxmc-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Jamie Cheshire
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.