Merge final changes before release#121
Merged
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughUpdates the README badges, adds shared scalar typing aliases, and propagates them through core single-event helpers, likelihoods, transforms, and sampler contracts. Several sampler cleanup paths also remove postponed annotation imports and clear JAX compilation cache settings. ChangesScalar typing and cleanup
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
After shutil.rmtree removes the jax_cache directory, the global jax_compilation_cache_dir config was left pointing at the deleted path. Post-sampling JAX JIT calls would then attempt to write to that path, causing FileNotFoundError on cache writes and breaking flowMC's `if current_cache is None` guard for any Sampler constructed later in the same process. Reset to None (the field's default) rather than "" — None correctly signals "no cache configured" to both JAX internals and flowMC's guard. Also adds the missing reset to _run_adaptive_persistent in SMC (the default AP mode was the only _run_* method that had rmtree without the config reset). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Member
Author
|
@coderabbitai review |
Contributor
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/jimgw/core/single_event/likelihood.py`:
- Line 403: Restore the Ruff suppression on the jaxtyping annotation in
_reduce_time so the type hint does not trigger F722; add the same # noqa: F722
comment style used elsewhere in the codebase (for example in flowmc.py) to the
Float[Array, " n_freq"] annotation in the likelihood.py single_event likelihood
helper.
In `@src/jimgw/samplers/blackjax/ns_aw.py`:
- Around line 256-258: The JAX cache reset in the cleanup path only runs on the
successful exit path, so a failure after configure_jax_cache() can leave
jax.config’s compilation cache directory pointing at the checkpoint tree. Move
the cache cleanup into a failure-safe finalization path in ns_aw.py, using the
existing config.checkpoint_dir handling and
jax.config.update("jax_compilation_cache_dir", None) so it executes even when
sampling or checkpoint finalisation raises. Keep the cache directory removal and
config reset tied to the same teardown logic regardless of success or failure.
In `@src/jimgw/samplers/blackjax/nss.py`:
- Around line 242-244: Move the JAX cache cleanup in `NSS` sampling into a
`finally` block so it always runs even if `_sample` raises. Update the control
flow around `configure_jax_cache()` and `_sample` in `blackjax/nss.py` so
`shutil.rmtree(config.checkpoint_dir / "jax_cache", ignore_errors=True)` and
`jax.config.update("jax_compilation_cache_dir", None)` are guaranteed to execute
after the sampling attempt, regardless of success or failure.
In `@src/jimgw/samplers/blackjax/smc.py`:
- Around line 274-276: The JAX cache reset in the runner cleanup paths is only
happening when checkpoint_dir is set on the happy path, so failures can leave
jax_compilation_cache_dir pointing at a stale checkpoint location for the rest
of the process. Update the cleanup logic around configure_jax_cache() in the SMC
runner flows (including the matching cleanup blocks in the other runner modes)
so the global JAX cache config is always reset unconditionally, even when
sampling, checkpoint I/O, or post-processing fails; keep the cache directory
removal guarded by checkpoint_dir, but move the config reset into a guaranteed
cleanup/finally path.
In `@src/jimgw/samplers/flowmc.py`:
- Around line 98-106: The F722 suppressions in _logpdf_flowmc and
_logprior_flowmc are attached to the wrong physical lines, so Ruff still flags
the Float[Array, " n_dims"] annotations. Move the # noqa: F722 comments onto the
parameter-annotation lines in those two methods (or otherwise suppress the exact
offending lines), and keep the fix localized to _logpdf_flowmc and
_logprior_flowmc now that from __future__ import annotations is gone.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b35a34d1-45b6-4f6f-a026-81a7cfd9336a
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (24)
README.mdsrc/jimgw/core/base.pysrc/jimgw/core/jim.pysrc/jimgw/core/prior.pysrc/jimgw/core/single_event/data.pysrc/jimgw/core/single_event/detector.pysrc/jimgw/core/single_event/likelihood.pysrc/jimgw/core/single_event/polarization.pysrc/jimgw/core/single_event/time_utils.pysrc/jimgw/core/single_event/transform_utils.pysrc/jimgw/core/single_event/transforms.pysrc/jimgw/core/single_event/utils.pysrc/jimgw/core/transforms.pysrc/jimgw/samplers/__init__.pysrc/jimgw/samplers/base.pysrc/jimgw/samplers/blackjax/_acceptance_walk_kernel.pysrc/jimgw/samplers/blackjax/_imports.pysrc/jimgw/samplers/blackjax/ns_aw.pysrc/jimgw/samplers/blackjax/nss.pysrc/jimgw/samplers/blackjax/smc.pysrc/jimgw/samplers/config.pysrc/jimgw/samplers/flowmc.pysrc/jimgw/samplers/periodic.pysrc/jimgw/typing.py
💤 Files with no reviewable changes (4)
- src/jimgw/samplers/init.py
- src/jimgw/samplers/periodic.py
- src/jimgw/samplers/blackjax/_acceptance_walk_kernel.py
- src/jimgw/samplers/blackjax/_imports.py
Comment on lines
256
to
+258
| if config.checkpoint_dir is not None: | ||
| shutil.rmtree(config.checkpoint_dir / "jax_cache", ignore_errors=True) | ||
| jax.config.update("jax_compilation_cache_dir", None) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Make the JAX cache reset run on failure as well.
The new reset on Line 258 only executes after a clean exit. If sampling or checkpoint finalisation raises after configure_jax_cache(), the process-global jax_compilation_cache_dir stays pinned to this checkpoint tree and the stale-cache-path issue still leaks into the next run.
Suggested fix
- config.configure_jax_cache()
- _method_t0 = time.perf_counter()
-
- def _validated_initial_particles(pos):
- ...
-
- ...
-
- if ckpt_path is not None:
- ckpt_path.unlink(missing_ok=True)
- if config.checkpoint_dir is not None:
- shutil.rmtree(config.checkpoint_dir / "jax_cache", ignore_errors=True)
- jax.config.update("jax_compilation_cache_dir", None)
+ config.configure_jax_cache()
+ _method_t0 = time.perf_counter()
+ try:
+ def _validated_initial_particles(pos):
+ ...
+
+ ...
+
+ if ckpt_path is not None:
+ ckpt_path.unlink(missing_ok=True)
+ finally:
+ if config.checkpoint_dir is not None:
+ shutil.rmtree(config.checkpoint_dir / "jax_cache", ignore_errors=True)
+ jax.config.update("jax_compilation_cache_dir", None)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if config.checkpoint_dir is not None: | |
| shutil.rmtree(config.checkpoint_dir / "jax_cache", ignore_errors=True) | |
| jax.config.update("jax_compilation_cache_dir", None) | |
| config.configure_jax_cache() | |
| _method_t0 = time.perf_counter() | |
| try: | |
| def _validated_initial_particles(pos): | |
| ... | |
| ... | |
| if ckpt_path is not None: | |
| ckpt_path.unlink(missing_ok=True) | |
| finally: | |
| if config.checkpoint_dir is not None: | |
| shutil.rmtree(config.checkpoint_dir / "jax_cache", ignore_errors=True) | |
| jax.config.update("jax_compilation_cache_dir", None) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/jimgw/samplers/blackjax/ns_aw.py` around lines 256 - 258, The JAX cache
reset in the cleanup path only runs on the successful exit path, so a failure
after configure_jax_cache() can leave jax.config’s compilation cache directory
pointing at the checkpoint tree. Move the cache cleanup into a failure-safe
finalization path in ns_aw.py, using the existing config.checkpoint_dir handling
and jax.config.update("jax_compilation_cache_dir", None) so it executes even
when sampling or checkpoint finalisation raises. Keep the cache directory
removal and config reset tied to the same teardown logic regardless of success
or failure.
Comment on lines
242
to
+244
| if config.checkpoint_dir is not None: | ||
| shutil.rmtree(config.checkpoint_dir / "jax_cache", ignore_errors=True) | ||
| jax.config.update("jax_compilation_cache_dir", None) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Move the cache teardown into a finally block.
Lines 242-244 only run when _sample completes normally. Any exception after configure_jax_cache() leaves jax_compilation_cache_dir pointing at this checkpoint location, so a failed NSS run can still poison later JAX work in the same process.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/jimgw/samplers/blackjax/nss.py` around lines 242 - 244, Move the JAX
cache cleanup in `NSS` sampling into a `finally` block so it always runs even if
`_sample` raises. Update the control flow around `configure_jax_cache()` and
`_sample` in `blackjax/nss.py` so `shutil.rmtree(config.checkpoint_dir /
"jax_cache", ignore_errors=True)` and
`jax.config.update("jax_compilation_cache_dir", None)` are guaranteed to execute
after the sampling attempt, regardless of success or failure.
Comment on lines
274
to
+276
| if config.checkpoint_dir is not None: | ||
| shutil.rmtree(config.checkpoint_dir / "jax_cache", ignore_errors=True) | ||
| jax.config.update("jax_compilation_cache_dir", None) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reset the global JAX cache config unconditionally in every runner.
These four cleanups still sit on the happy path only. Each mode calls configure_jax_cache() first, so any failure during sampling, checkpoint I/O, or post-processing leaves jax_compilation_cache_dir set to a stale checkpoint path for the rest of the process.
Also applies to: 395-397, 523-525, 653-655
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/jimgw/samplers/blackjax/smc.py` around lines 274 - 276, The JAX cache
reset in the runner cleanup paths is only happening when checkpoint_dir is set
on the happy path, so failures can leave jax_compilation_cache_dir pointing at a
stale checkpoint location for the rest of the process. Update the cleanup logic
around configure_jax_cache() in the SMC runner flows (including the matching
cleanup blocks in the other runner modes) so the global JAX cache config is
always reset unconditionally, even when sampling, checkpoint I/O, or
post-processing fails; keep the cache directory removal guarded by
checkpoint_dir, but move the config reset into a guaranteed cleanup/finally
path.
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.
Summary by CodeRabbit
Documentation
Refactor
Bug Fixes