Validate num_samples in MCMC initialization - #2236
Conversation
Qazalbash
left a comment
Qazalbash
left a comment
There was a problem hiding this comment.
Thanks for tackling this — an early, explicit error beats an IndexError from deep inside fori_collect. I ran the cases below against this branch to check the guard's coverage, and I think it needs a bit of widening before it lands. Findings are ordered by severity.
1. The guard checks the value but not the type
num_samples < 1 accepts any non-integer that compares against 1, so the argument types users actually get wrong still slip through:
MCMC(NUTS(model), num_warmup=2, num_samples=10.0) # constructs fine
# .run() -> TypeError: Shapes must be 1D sequences of concrete values of integer type, got (10.0,)
MCMC(NUTS(model), num_warmup=2, num_samples=2.5) # constructs fine
# silently truncates to 2 samples, no warningnum_samples=True is also accepted as 1 (bool subclasses int). The thinning check two lines below already has the right shape for this:
if not isinstance(thinning, int) or thinning < 1:The 2.5 case is the one I'd most want to catch — it produces a quietly wrong number of draws rather than an error.
2. TypeError instead of ValueError for None
MCMC(NUTS(model), num_warmup=2, num_samples=None)
# TypeError: '<' not supported between instances of 'NoneType' and 'int'None is the common case when num_samples comes from a config file or CLI arg with a missing key, and callers wrapping construction in except ValueError won't catch it. The isinstance form in point 1 fixes this for free.
3. This rejects a warmup-only workflow that currently works
MCMC(..., num_samples=0) followed by mcmc.warmup(key, collect_warmup=True) works today and is now rejected at construction. Verified on this branch with the guard removed:
mcmc = MCMC(NUTS(model), num_warmup=6, num_samples=0)
mcmc.warmup(random.key(0), collect_warmup=True)
mcmc.get_samples() # {'x': (6,)} -- correctThe reason it works is that the collect_warmup=True branch never reads self.num_samples:
Lines 626 to 631 in 88ad9fb
It passes self.num_warmup as the collection size. So adaptation-only runs (tune step size / mass matrix, stash post_warmup_state, sample later) are a legitimate use of num_samples=0. The collect_warmup=False path genuinely is broken with 0, so the constraint is real — it just isn't a constructor-level invariant. Validating in run() / the collect_warmup=False branch would catch the actual breakage without rejecting the working case.
4. num_samples >= 1 isn't the real precondition — thinning interacts
MCMC(NUTS(model), num_warmup=4, num_samples=1, thinning=2)
# passes both guards, then .run() -> IndexError: index is out of bounds for axis 0 with size 0collection_size = collection_size // self.thinning gives 1 // 2 == 0:
Lines 493 to 497 in 88ad9fb
Since this PR adds validation immediately next to the thinning check, num_samples >= thinning seems like the constraint worth enforcing.
5. num_warmup on the line above has a worse, silent failure
mcmc = MCMC(NUTS(model), num_warmup=-5, num_samples=4)
mcmc.run(random.key(0))
mcmc.get_samples()['x'] # [-0. -0. -0. -0.]No error at all — uninitialized buffer contents returned as posterior draws. _set_collection_params computes lower=-5, upper=-1, fori_collect's assert lower <= upper passes, and the loop body never runs. num_chains=0 is also unvalidated (IndexError: tuple index out of range).
Silently wrong numbers are worse than the crash this PR fixes, and it's the adjacent line. Rather than three more one-off if statements, a small shared helper validating num_warmup, num_samples, num_chains, and thinning uniformly would fix the whole family at the right depth.
Smaller points
- No test. The new branch has no coverage, so a future refactor could drop it silently. (The neighbouring
thinningcheck is untested too, so there's no existing pattern to follow — a smallpytest.raises(ValueError)parametrized over0and-1intest/infer/test_mcmc.pywould cover both.) - Docstring. Line 272 still reads
:param int num_samples: Number of samples to generate from the Markov chain.The constraint isn't discoverable except by triggering it; comparethinning, documented asPositive integer that controls.... - Validation only at
__init__.num_samplesis a plain public attribute read live at run time (_set_collection_params,_compile,warmup,run), somcmc.num_samples = 0after construction reinstates every failure mode. Worth noting if the intent is a real invariant rather than a typo catcher. - Message wording.
"num_samples must be at least 1"vs. line 349's"thinning must be a positive integer"for the same class of constraint. Including the received value would also help when the argument is computed rather than literal.
Happy to be wrong on #3 if warmup-only with num_samples=0 isn't considered supported — but if it's being dropped deliberately, that's worth calling out as a behaviour change.
Review generated with Claude Code
Benchmark reportthis PR run time: unchanged across 32 benchmarks
compile time: unchanged across 32 benchmarksNo significant changesEvery benchmark stayed within ±5% run time and ±25% compile time. Red is slower, green is faster; a row is coloured by the worse of its two columns. A delta in parentheses cleared the threshold on a measurement below the resolution floor, so it is shown without being called a change. † marks a benchmark that could not be compared — see below. Full results
|
| baseline | this PR | |
|---|---|---|
| ref | master |
patch-1 |
| commit | a47333ee |
182a5ac3 |
| numpyro | 0.21.0 | 0.21.0 |
| jax | 0.11.0 | 0.11.0 |
| backend | cpu | cpu |
| python | 3.14.7 | 3.14.7 |
Runner: Linux-6.17.0-1020-azure-x86_64-with-glibc2.39, 4 CPUs.
Produced by this benchmark run.
2 participants
I tried to debug my MCMC kernel by doing just one iteration using
I got this very long error message
and spend a lot of time debugging my kernel for indexing mistakes.
This PR adds a check for
num_samplesthat hopefully saves the next person from doing that.