diff --git a/CHANGELOG.md b/CHANGELOG.md index f7a850e..37e0ef4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,15 @@ ## 0.16.0 — unreleased +- **`NUTS`** — the No-U-Turn Sampler (Hoffman & Gelman 2014, multinomial variant): + HMC that tunes its own trajectory length (doubling until a whole-span U-turn) and + step size (dual-averaging warmup to `target_accept`), so it needs neither + `leapfrog_steps` nor a step-size sweep. All chains are evolved in lockstep with a + per-chain freeze mask, so each chain's draw is identical to independent + single-chain NUTS. Exposes `last_tree_depth` and a `divergences` count. Validated: + recovers a standard and a correlated Gaussian's covariance, tunes to 0.8 + acceptance, enters Neal's funnel, and terminates by U-turn (small tree depths, + zero divergences on the Gaussian). Example `nuts_sampling.py`. - **`LatentEBM`** — a latent-variable EBM: a joint `E(x, z) = E_prior(z) + E(x | z)` coupling a prior over a latent `z` (default standard normal) with a decoder energy. The data marginal is intractable, so `sample_joint` runs block diff --git a/README.md b/README.md index 649db4c..c0aac6c 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,7 @@ returns `LossOutput(loss, metrics, x_neg)`; call `out.loss.backward()`). | Piece | Contents | |---|---| | **Energies** | any callable `(B, *shape) -> (B,)`; `nets.MLPEnergy` / `ConvEnergy` / `ResNetEnergy` / `ConvClassifier` (SiLU, optional spectral norm, no batch norm), `nets.RBM` (Bernoulli RBM with exact `log_z`), `GuidedEnergy` (classifier-free guidance), `nets.IsingEnergy` / `PottsEnergy` (discrete lattices), `nets.FunnelEnergy` / `GaussianMixtureEnergy` / `BananaEnergy` (closed-form targets), `nets.AffineCouplingFlow` (RealNVP) / `nets.NeuralSplineCouplingFlow` (rational-quadratic spline) / `nets.ContinuousNormalizingFlow` (FFJORD — trainable exact-likelihood flows / self-normalized energies), noise-conditional variants for NCSN; `EnergyModel`, `ebm.score` | -| **Samplers** | `LangevinDynamics` (ULA/SGLD), `MALA`, `AdaptiveMALA` (dual-averaging step-size warmup + diagonal metric), `HMC`, `UnderdampedLangevin` (SGHMC), `PreconditionedLangevin`, `ParallelTempering` (replica exchange), `TemperedTransitions`, `SVGD` (Stein variational), `GibbsSampler` (block Gibbs), `GibbsWithGradients` + `CategoricalGibbsWithGradients`, `AnnealedLangevinDynamics`, `ProbabilityFlowODE` / `PredictorCorrector` (score-SDE), `DDPMAncestralSampler` (VP diffusion) | +| **Samplers** | `LangevinDynamics` (ULA/SGLD), `MALA`, `AdaptiveMALA` (dual-averaging step-size warmup + diagonal metric), `HMC`, `NUTS` (No-U-Turn Sampler, self-tuning trajectory length), `UnderdampedLangevin` (SGHMC), `PreconditionedLangevin`, `ParallelTempering` (replica exchange), `TemperedTransitions`, `SVGD` (Stein variational), `GibbsSampler` (block Gibbs), `GibbsWithGradients` + `CategoricalGibbsWithGradients`, `AnnealedLangevinDynamics`, `ProbabilityFlowODE` / `PredictorCorrector` (score-SDE), `DDPMAncestralSampler` (VP diffusion) | | **Losses** | `ContrastiveDivergence` (CD-k / persistent CD), `DiffusionRecoveryLikelihood` + `drl_sample`, `DenoisingScoreMatching` / `MultiSigmaDenoisingScoreMatching` (NCSN), `VPDenoisingScoreMatching` (DDPM), `SlicedScoreMatching`, `ExactScoreMatching`, `EnergyDiscrepancy` (MCMC-free), `PseudoLikelihood` / `RatioMatching` / `ConcreteScoreMatching` (MCMC-free, discrete), `NoiseContrastiveEstimation`, `JEMLoss` | | **Composition** | `SumEnergy` (product of experts), `MixtureEnergy`, `EnsembleEnergy` (deep-ensemble mean energy + member disagreement), `TemperedEnergy`, `LatentEBM` (joint `E(x, z)` with a prior + decoder, block-Gibbs sampled) — energies compose like densities and nest | | **Training** | thin `Trainer` (device, EMA, supervised batches, `save`/`load` checkpointing), `ReplayBuffer`, `EMA` | @@ -76,6 +76,7 @@ Runnable scripts in [`examples/`](https://github.com/davidkhjo/ebmkit/tree/main/ - `train_energy_discrepancy.py` — two-moons trained MCMC-free (energy discrepancy) - `sampling_hard_targets.py` — parallel tempering escapes a trapped mode; ESS / R̂ diagnostics - `adaptive_mala.py` — a self-tuning MALA: dual-averaging step size + a learned diagonal metric +- `nuts_sampling.py` — the No-U-Turn Sampler on Neal's funnel: trajectory length adapts per draw - `goodness_of_fit.py` — KSD for model selection; classifier two-sample test - `ensemble_ood.py` — a deep-ensemble EBM whose member disagreement flags OOD - `latent_ebm.py` — a latent-variable EBM: block-Gibbs on a joint `E(x, z)` matches ancestral sampling diff --git a/examples/nuts_sampling.py b/examples/nuts_sampling.py new file mode 100644 index 0000000..71c37a4 --- /dev/null +++ b/examples/nuts_sampling.py @@ -0,0 +1,70 @@ +"""NUTS: HMC that picks its own trajectory length, on Neal's funnel. + +The No-U-Turn Sampler removes HMC's two hand-tuned knobs: it doubles each +trajectory until the path starts to double back (a U-turn), and it tunes the step +size to a target acceptance by dual averaging during a warmup. No `leapfrog_steps`, +no step-size sweep. Neal's funnel — a Gaussian whose width is itself a Gaussian +latent `v` — is the classic stress test: the neck is sharp where `v` is negative, +so the sampler must *lengthen* its trajectories there. We plot the samples and the +distribution of tree depths (how far NUTS doubled each draw), and print the tuned +step size, mean acceptance, and divergence count. + +Run: python examples/nuts_sampling.py +Outputs nuts_sampling_result.png next to this script (needs the [viz] extra). +""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +import ebm + + +def main() -> None: + torch.manual_seed(0) + energy = ebm.nets.FunnelEnergy(dim=2, v_scale=3.0) # v ~ N(0, 9), x | v ~ N(0, e^v) + + sampler = ebm.NUTS(step_size=0.3, steps=300, warmup=300, target_accept=0.8) + x = sampler.sample(energy, torch.randn(3000, 2)) + print(f"tuned step size = {sampler.step_size:.3f}") + print(f"mean acceptance = {sampler.last_accept_rate:.3f} (target 0.8)") + print( + f"v marginal std = {x[:, 0].std():.2f} (true 3.0; identity metric under-samples the neck)" + ) + print(f"divergences = {sampler.divergences}") + + # Collect tree depths over a batch of post-warmup draws (step size now frozen). + depths = [] + xd = x + for _ in range(40): + xd = sampler.step(energy, xd) + depths.append(sampler.last_tree_depth) + depths = torch.stack(depths).reshape(-1) + + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, axes = plt.subplots(1, 2, figsize=(12, 5)) + axes[0].scatter(x[:, 1], x[:, 0], s=5, alpha=0.3, color="#5c50c9") + axes[0].set_xlabel("x (neck coordinate)") + axes[0].set_ylabel("v (log-scale latent)") + axes[0].set_title("NUTS samples of Neal's funnel") + axes[0].set_xlim(-15, 15) + + hi = int(depths.max().item()) + axes[1].hist(depths.numpy(), bins=range(hi + 2), align="left", rwidth=0.8, color="#5c50c9") + axes[1].set_xlabel("tree depth reached") + axes[1].set_ylabel("draws") + axes[1].set_title("Trajectory length adapts per draw (U-turn, not a fixed length)") + + out = Path(__file__).parent / "nuts_sampling_result.png" + fig.savefig(out, dpi=120, bbox_inches="tight") + print(f"saved {out}") + + +if __name__ == "__main__": + main() diff --git a/examples/nuts_sampling_result.png b/examples/nuts_sampling_result.png new file mode 100644 index 0000000..7d8e346 Binary files /dev/null and b/examples/nuts_sampling_result.png differ diff --git a/src/ebm/__init__.py b/src/ebm/__init__.py index 816cb77..92b3fbd 100644 --- a/src/ebm/__init__.py +++ b/src/ebm/__init__.py @@ -31,6 +31,7 @@ from ebm.samplers import ( HMC, MALA, + NUTS, SVGD, AdaptiveMALA, AnnealedLangevinDynamics, @@ -55,6 +56,7 @@ "EMA", "HMC", "MALA", + "NUTS", "SVGD", "AISResult", "AdaptiveMALA", diff --git a/src/ebm/samplers/__init__.py b/src/ebm/samplers/__init__.py index 85ca709..0ef92ae 100644 --- a/src/ebm/samplers/__init__.py +++ b/src/ebm/samplers/__init__.py @@ -10,6 +10,7 @@ PreconditionedLangevin, UnderdampedLangevin, ) +from ebm.samplers.nuts import NUTS from ebm.samplers.score_sde import PredictorCorrector, ProbabilityFlowODE from ebm.samplers.svgd import SVGD from ebm.samplers.tempering import ParallelTempering, TemperedTransitions @@ -17,6 +18,7 @@ __all__ = [ "HMC", "MALA", + "NUTS", "SVGD", "AdaptiveMALA", "AnnealedLangevinDynamics", diff --git a/src/ebm/samplers/nuts.py b/src/ebm/samplers/nuts.py new file mode 100644 index 0000000..56af096 --- /dev/null +++ b/src/ebm/samplers/nuts.py @@ -0,0 +1,289 @@ +"""No-U-Turn Sampler (NUTS): HMC that picks its own trajectory length.""" + +from __future__ import annotations + +import math + +import torch +from torch import Tensor, nn + +from ebm._functional import flat_sum as _flat_sum +from ebm.energy import EnergyFn +from ebm.samplers.base import Sampler +from ebm.utils import frozen_params + +_NEG_INF = float("-inf") + + +def _bcast(mask: Tensor, like: Tensor) -> Tensor: + """Reshape a ``(B,)`` mask to broadcast against ``like``'s event dims.""" + return mask.reshape(-1, *([1] * (like.dim() - 1))) + + +class NUTS(Sampler): + """No-U-Turn Sampler (Hoffman & Gelman 2014), the multinomial variant. + + HMC with two things automated: the trajectory length (doubled until the path + makes a U-turn, so it needs no ``leapfrog_steps``) and the step size (tuned to + ``target_accept`` by dual averaging during a warmup, then frozen). Best on + smooth, low-dimensional continuous targets; unbiased (targets exactly + ``p ∝ exp(-E)``). + + All ``B`` chains are evolved together in **lockstep**: they build their trees to + a shared depth each doubling, but a chain that has already met its stop + criterion (a whole-span U-turn or a divergence) is frozen — every state write + is masked, so a stopped chain's draw is identical to what an independent + single-chain NUTS would have produced. Wasted leapfrogs on finished chains are + the only cost; correctness is unaffected. + + Diagnostics after a run: ``last_accept_rate`` (mean Metropolis acceptance of + the last draw), ``last_tree_depth`` (per-chain tree depth reached), and + ``divergences`` (count over the sampling phase — persistently nonzero means the + step size is too large or the geometry too sharp for an identity metric, e.g. + Neal's funnel). + + Args: + step_size: initial ε (a starting guess; warmup overwrites it). + steps: default number of post-warmup draws per ``sample`` call. + warmup: dual-averaging iterations (0 disables adaptation). + target_accept: δ the warmup targets (0.8 is the NUTS default). + max_depth: cap on tree depth (``2**max_depth`` leapfrogs per draw). + max_delta_h: divergence threshold on the Hamiltonian error. + gamma, t0, kappa: dual-averaging shrinkage / stabilization / decay constants. + """ + + def __init__( + self, + step_size: float = 0.1, + steps: int = 100, + *, + warmup: int = 1000, + target_accept: float = 0.8, + max_depth: int = 10, + max_delta_h: float = 1000.0, + gamma: float = 0.05, + t0: float = 10.0, + kappa: float = 0.75, + ): + super().__init__(steps) + if not 0.0 < target_accept < 1.0: + raise ValueError("target_accept must be in (0, 1)") + if warmup < 0: + raise ValueError("warmup must be >= 0") + if max_depth < 1: + raise ValueError("max_depth must be >= 1") + self.step_size = step_size + self.warmup = warmup + self.target_accept = target_accept + self.max_depth = max_depth + self.max_delta_h = max_delta_h + self.gamma = gamma + self.t0 = t0 + self.kappa = kappa + self.last_tree_depth: Tensor | None = None + self.divergences = 0 + + def _leapfrog( + self, energy: EnergyFn, x: Tensor, p: Tensor, grad: Tensor, signed_eps: Tensor + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """One leapfrog step with per-chain signed step ``signed_eps`` (``v·ε``).""" + e = signed_eps.reshape(-1, *([1] * (x.dim() - 1))) + p = p - 0.5 * e * grad + x = x + e * p + e_new, grad = self._energy_grad(energy, x) + p = p - 0.5 * e * grad + return x, p, grad, e_new + + @staticmethod + def _no_uturn(x_minus: Tensor, x_plus: Tensor, p: Tensor) -> Tensor: + return _flat_sum((x_plus - x_minus) * p) >= 0 + + def _span_ok(self, xm: Tensor, xp: Tensor, pm: Tensor, pp: Tensor) -> Tensor: + """Whole-span no-U-turn: the span still advances at both ends.""" + return self._no_uturn(xm, xp, pm) & self._no_uturn(xm, xp, pp) + + def _build_tree( + self, + energy: EnergyFn, + x: Tensor, + p: Tensor, + grad: Tensor, + v: Tensor, + depth: int, + eps: float, + h0: Tensor, + active: Tensor, + ) -> tuple[Tensor, ...]: + """Recursively double the trajectory; every write gated by ``active``. + + Returns ``(x⁻, p⁻, g⁻, x⁺, p⁺, g⁺, x_prop, logw, s, diverged, a_sum, n_a)`` + with log-space multinomial weight ``logw = logsumexp(H0 − H)`` over the + subtree's leaves and validity mask ``s``. + """ + if depth == 0: + x1, p1, g1, e1 = self._leapfrog(energy, x, p, grad, v * eps) + h1 = e1 + 0.5 * _flat_sum(p1.pow(2)) + d_h = h0 - h1 + finite = torch.isfinite(h1) + over = (h1 - h0) > self.max_delta_h + diverged = active & (~finite | over) + valid = active & finite & ~over + m = _bcast(active, x1) # freeze inactive chains at their input state + x1 = torch.where(m, x1, x) + p1 = torch.where(m, p1, p) + g1 = torch.where(m, g1, grad) + logw = torch.where(valid, d_h, torch.full_like(d_h, _NEG_INF)) + a = torch.where(active, torch.exp(d_h.clamp(max=0.0)), torch.zeros_like(d_h)) + n_a = active.to(d_h.dtype) + return (x1, p1, g1, x1, p1, g1, x1, logw, valid, diverged, a, n_a) + + xm, pm, gm, xp, pp, gp, prop1, logw1, s1, d1, a1, na1 = self._build_tree( + energy, x, p, grad, v, depth - 1, eps, h0, active + ) + active2 = active & s1 # only chains still valid extend a second subtree + plus = _bcast(v > 0, xm) + x2s = torch.where(plus, xp, xm) + p2s = torch.where(plus, pp, pm) + g2s = torch.where(plus, gp, gm) + xm2, pm2, gm2, xp2, pp2, gp2, prop2, logw2, s2, d2, a2, na2 = self._build_tree( + energy, x2s, p2s, g2s, v, depth - 1, eps, h0, active2 + ) + new_xm = torch.where(plus, xm, xm2) + new_pm = torch.where(plus, pm, pm2) + new_gm = torch.where(plus, gm, gm2) + new_xp = torch.where(plus, xp2, xp) + new_pp = torch.where(plus, pp2, pp) + new_gp = torch.where(plus, gp2, gp) + + denom = torch.logaddexp(logw1, logw2) + log_u = torch.log(torch.rand_like(logw1)) + replace = (log_u < (logw2 - denom)) & s2 & active2 # multinomial pick, -inf-safe + prop = torch.where(_bcast(replace, prop1), prop2, prop1) + + no_uturn = self._span_ok(new_xm, new_xp, new_pm, new_pp) + s = s1 & s2 & no_uturn & active + return ( + new_xm, + new_pm, + new_gm, + new_xp, + new_pp, + new_gp, + prop, + denom, + s, + d1 | d2, + a1 + a2, + na1 + na2, + ) + + def _draw( + self, energy: EnergyFn, x0: Tensor, eps: float + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """One NUTS transition for the whole batch; returns ``(x, ᾱ, depth, diverged)``.""" + b = x0.shape[0] + p = torch.randn_like(x0) + e0, grad0 = self._energy_grad(energy, x0) + h0 = e0 + 0.5 * _flat_sum(p.pow(2)) + + xm = xp = x0 + pm = pp = p + gm = gp = grad0 + x_prop = x0 + logw = torch.zeros(b, device=x0.device, dtype=x0.dtype) + alive = torch.ones(b, device=x0.device, dtype=torch.bool) + a_tot = torch.zeros(b, device=x0.device, dtype=x0.dtype) + na_tot = torch.zeros(b, device=x0.device, dtype=x0.dtype) + depth_reached = torch.zeros(b, device=x0.device, dtype=torch.long) + diverged = torch.zeros(b, device=x0.device, dtype=torch.bool) + + for depth in range(self.max_depth): + if not bool(alive.any()): + break + depth_reached += alive.long() + v = torch.where( + torch.rand(b, device=x0.device) < 0.5, + torch.full((b,), -1.0, device=x0.device), + torch.full((b,), 1.0, device=x0.device), + ) + plus = _bcast(v > 0, xm) + xs = torch.where(plus, xp, xm) + ps = torch.where(plus, pp, pm) + gs = torch.where(plus, gp, gm) + nm, npm, ngm, npx, npp, ngp, prop_s, logw_s, s_s, d_s, a_s, na_s = self._build_tree( + energy, xs, ps, gs, v, depth, eps, h0, alive + ) + upd_p = _bcast(alive & (v > 0), xm) + upd_m = _bcast(alive & (v < 0), xm) + xp = torch.where(upd_p, npx, xp) + pp = torch.where(upd_p, npp, pp) + gp = torch.where(upd_p, ngp, gp) + xm = torch.where(upd_m, nm, xm) + pm = torch.where(upd_m, npm, pm) + gm = torch.where(upd_m, ngm, gm) + + log_u = torch.log(torch.rand_like(logw)) + replace = (log_u < (logw_s - logw)) & s_s & alive + x_prop = torch.where(_bcast(replace, x_prop), prop_s, x_prop) + logw = torch.logaddexp( + logw, torch.where(alive, logw_s, torch.full_like(logw_s, _NEG_INF)) + ) + a_tot += torch.where(alive, a_s, torch.zeros_like(a_s)) + na_tot += torch.where(alive, na_s, torch.zeros_like(na_s)) + diverged |= d_s & alive + alive = alive & s_s & self._span_ok(xm, xp, pm, pp) + + alpha_bar = a_tot / na_tot.clamp_min(1.0) + return x_prop.detach(), alpha_bar, depth_reached, diverged + + def step(self, energy: EnergyFn, x: Tensor) -> Tensor: + x_next, alpha_bar, depth, diverged = self._draw(energy, x.detach(), self.step_size) + self._last_accept = alpha_bar.mean() + self.last_tree_depth = depth + self.divergences += int(diverged.sum()) + return x_next + + def _dual_average( + self, energy: EnergyFn, x: Tensor, n: int, eps0: float + ) -> tuple[Tensor, float]: + """Run ``n`` dual-averaging warmup draws from ``eps0``; return ``(x, ε̄)``.""" + if n == 0: + return x, eps0 + mu = math.log(10 * eps0) + log_eps = math.log(eps0) + log_ebar = 0.0 + h_bar = 0.0 + for m in range(1, n + 1): + x, alpha_bar, _, _ = self._draw(energy, x, math.exp(log_eps)) + x = x.detach() + gap = self.target_accept - float(alpha_bar.mean()) + h_bar = (1 - 1 / (m + self.t0)) * h_bar + gap / (m + self.t0) + log_eps = mu - math.sqrt(m) / self.gamma * h_bar + eta = m**-self.kappa + log_ebar = eta * log_eps + (1 - eta) * log_ebar + return x, math.exp(log_ebar) + + def sample( + self, + energy: EnergyFn, + x_init: Tensor, + *, + steps: int | None = None, + return_trajectory: bool = False, + ) -> Tensor: + """Warm up (tuning ε), freeze, then draw. See the class docstring.""" + n_steps = self.steps if steps is None else steps + x = x_init.detach().clone() + self.divergences = 0 + module = energy if isinstance(energy, nn.Module) else None + with frozen_params(module), torch.enable_grad(): + x, eps = self._dual_average(energy, x, self.warmup, self.step_size) + self.step_size = eps # freeze the averaged step size + trajectory = [x.clone()] if return_trajectory else None + for _ in range(n_steps): + x = self.step(energy, x).detach() + if trajectory is not None: + trajectory.append(x.clone()) + if trajectory is not None: + return torch.stack(trajectory) + return x diff --git a/tests/test_samplers.py b/tests/test_samplers.py index 757bc85..cfaed93 100644 --- a/tests/test_samplers.py +++ b/tests/test_samplers.py @@ -245,3 +245,60 @@ def test_adaptive_mala_zero_warmup_keeps_step_size_and_validates(): ebm.AdaptiveMALA(target_accept=1.5) with pytest.raises(ValueError): ebm.AdaptiveMALA(warmup=-1) + + +def test_nuts_targets_standard_normal_and_tunes_step_size(): + sampler = ebm.NUTS(step_size=0.5, steps=80, warmup=200) + samples = sampler.sample(quadratic_energy, 3 * torch.randn(2000, 2)) + _check_standard_normal(samples) + assert not samples.requires_grad + assert abs(sampler.last_accept_rate - 0.8) < 0.1 # dual averaging hits the target + assert sampler.divergences == 0 # a smooth Gaussian never diverges + + +def test_nuts_recovers_a_correlated_gaussian(): + cov = torch.tensor([[2.0, 1.2], [1.2, 1.5]]) + sampler = ebm.NUTS(step_size=0.5, steps=160, warmup=250) + traj = sampler.sample( + _correlated_gaussian(cov), 3 * torch.randn(3000, 2), return_trajectory=True + ) + pooled = traj[-20:].reshape(-1, 2) # pool near-independent late draws to cut MC noise + assert (torch.cov(pooled.T) - cov).abs().max().item() < 0.15 + + +def test_nuts_explores_neals_funnel(): + # Identity-metric NUTS enters but under-samples the funnel neck, so v.std biases + # below the true 3 — assert only that it stays finite and spreads (loose band). + energy = ebm.nets.FunnelEnergy(dim=2, v_scale=3.0) + sampler = ebm.NUTS(step_size=0.3, steps=150, warmup=200) + samples = sampler.sample(energy, torch.randn(1500, 2)) + assert torch.isfinite(samples).all() + assert 1.8 < samples[:, 0].std().item() < 3.3 + + +def test_nuts_tree_depth_terminates_by_uturn(): + # On an easy Gaussian the U-turn fires at small depths — the tree should never + # bottom out at max_depth (which would mean it never detected a U-turn). + sampler = ebm.NUTS(step_size=0.5, steps=1, warmup=150) + x = sampler.sample(quadratic_energy, torch.randn(1024, 2)) + depths = [] + for _ in range(30): + x = sampler.step(quadratic_energy, x) + depths.append(sampler.last_tree_depth) + depths = torch.stack(depths) + assert (depths < sampler.max_depth).all() + assert depths.float().median().item() <= 3 + assert (depths <= 3).float().mean().item() > 0.8 + + +def test_nuts_accept_lifecycle_and_validation(): + sampler = ebm.NUTS(step_size=0.5, steps=5, warmup=0) + assert sampler.last_accept_rate is None # nothing drawn yet + sampler.sample(quadratic_energy, torch.randn(64, 2)) + assert isinstance(sampler.last_accept_rate, float) + with pytest.raises(ValueError): + ebm.NUTS(target_accept=1.5) + with pytest.raises(ValueError): + ebm.NUTS(warmup=-1) + with pytest.raises(ValueError): + ebm.NUTS(max_depth=0)