Expose MH acceptance rate and fix truncated normal tails - #188
Merged
Conversation
gen_params is jitted and its body assumed every scalar argument arrives as a traced array. jax does not trace an argument left at its default, so an omitted offset reached check_offset as a plain Python float and raised AttributeError. Convert it explicitly, in the same place the other scalars are cast, so the array also propagates to the returned Params. offset is the only argument with a non-None scalar default. gen_data was unaffected because it always forwards offset to the jitted gen_params, which traces it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`truncated_normal_onesided` could return ±inf: the lower target of the clip guard was the smallest subnormal float, which xla flushes to zero on cpu, and which ndtri maps to -inf anyway. Clip to the smallest normal float instead, which ndtri maps to ~-12.9. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Past |bound| ~ 12.9 in float32 the inversion of the normal cdf saturates, so the sample fell short of the truncation region; in probit this made the latent contradict the outcome it was conditioned on. Clamp the sample onto the bound, where the distribution concentrates for such bounds anyway. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Most tests in the module picked an arbitrary forest size (2, 3, 4, 5, 6, 10, 20). Route them all through a single NUM_TREES global so jax can reuse compiled code instead of recompiling per shape. The uv/mv equivalence check needs its rtol relaxed from 1e-6 to 1e-5, matching the error_cov_inv check next to it: with fewer trees the larger leaves push the float32 reduction noise just past the old threshold. test_blocked_mass_matches_reference keeps its own tree count, which is part of the tuning that makes the blocked mass nonzero. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The (n, p), (k, df) and k values were not picked for coverage: no library code branches on them in those ranges. The var dtype is uint8 up to p=256, prec_count_num_trees resolves to None for all these n, and the wishart and leaf term samplers are fully vmapped over k. Keep only the values with a role: k=1, which feeds the exactly zero matrix to chol_with_gersh; one wide and one tall data shape; and dof both tight and loose relative to k. Drop n=3, which left the mask in test_error_cov_inv_missing_equals_drop empty, and k=5, which equals NUM_TREES and so hides a trees/k axis swap. 290 -> 250 tests, 184s -> 163s. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`X` was a plain `arange`, so every row past the first was above the cutpoint range and sent all datapoints to the same child. The trees grew internal nodes but partitioned nothing: averaged over trees, all datapoints sat in a single leaf. Bin the index into the cutpoint range instead, rotating the phase per variable, which takes the occupied leaves per tree from 1.0 to 2.8 (n=10) and from 1.4 to 3.2 (n=50). Also assert that the mask in test_error_cov_inv_missing_equals_drop is neither empty nor full, since either way the two states it compares coincide and both assertions hold trivially. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add `Bart.accept` and `mc_gbart.accept`, the per-iteration fraction of trees with an accepted grow or prune move, matching the `accept` attribute of R BART3. Burn-in samples are included; unlike BART3, the iterations thinned away by `n_skip`/`keepevery` are not recorded. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds an
acceptproperty toBartandmc_gbartwith the per-iteration fraction of trees with an accepted grow/prune move, including burn-in and keeping chains separate to allow convergence checks.Alongside, it fixes
truncated_normal_onesidedfor bounds deep into a tail: the uniform to invert could underflow to 0 (especially on cpu, where xla flushes subnormals), producing infinities, and the inversion saturates at |x| ~ 12.9 in float32, which can leave the sample outside the truncation region. The sample is now clipped to the smallest normal number and collapsed onto the boundary when it falls short of it. This makesstep_zvalid when the latent is truncated far into a tail by an offset that contradicts the binary outcome.Minor:
gen_paramsnow converts the defaultoffsetto a jax array when called directly (outsidegen_datait was left an untraced Python float).Test suite:
test_mcmcstepshares oneNUM_TREESconstant and trims parametrized variants to cut recompilations, and its toy dataset is reworked so every variable is actually splittable.🤖 Generated with Claude Code