Batch take step#232
Conversation
|
Warning Rate limit exceeded@kazewong has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 6 minutes and 52 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (1)
WalkthroughThe changes introduce a new Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant RQSpline_MALA_Bundle
participant TakeSteps
User->>RQSpline_MALA_Bundle: Initialize (chain_batch_size=N)
RQSpline_MALA_Bundle->>TakeSteps: Initialize (chain_batch_size=N)
User->>TakeSteps: Call __call__ with chains
alt chain_batch_size batching
TakeSteps->>TakeSteps: Split chains into batches
loop For each batch
TakeSteps->>TakeSteps: Process batch with sample()
end
TakeSteps->>User: Concatenate and return results
else No batching
TakeSteps->>TakeSteps: Process all chains with sample()
TakeSteps->>User: Return results
end
Poem
✨ Finishing Touches
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
src/flowMC/resource/nf_model/NF_proposal.py (2)
15-19: Remove commented-out debug code.The commented-out debug imports and function should be removed to keep the code clean.
-# from jax import debug -# def cond_print(log_prob_proposal, log_prob_current, log_prob_nf_proposal, log_prob_nf_current, do_accept, position_proposal, position_current, tag=""): -# if (do_accept == True and ((log_prob_proposal - log_prob_current) < -5)): -# print(f"{tag} pro_prob: {log_prob_proposal}, cur_prob: {log_prob_current}, pro_nf_prob: {log_prob_nf_proposal}, cur_nf_prob: {log_prob_nf_current}, accept: {do_accept}, pro_pos: {position_proposal[0]}, cur_pos: {position_current[0]}") -
117-117: Remove commented-out debug callback.The commented-out debug callback should be removed to keep the code clean.
-# debug.callback(cond_print, log_prob_proposal, log_prob_initial, log_prob_nf_proposal, log_prob_nf_initial, do_accept, position_proposal, position_initial)src/flowMC/strategy/take_steps.py (1)
21-21: Fix inline comment formatting.The inline comment needs at least two spaces before it according to PEP 8.
- chain_batch_size: int # If vmap over a large number of chains is memory bounded, this splits the computation + chain_batch_size: int # If vmap over a large number of chains is memory bounded, this splits the computation
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
src/flowMC/resource/nf_model/NF_proposal.py(2 hunks)src/flowMC/resource_strategy_bundle/RQSpline_MALA.py(3 hunks)src/flowMC/resource_strategy_bundle/RQSpline_MALA_PT.py(3 hunks)src/flowMC/strategy/take_steps.py(4 hunks)
🧰 Additional context used
🪛 Flake8 (7.2.0)
src/flowMC/strategy/take_steps.py
[error] 21-21: at least two spaces before inline comment
(E261)
🪛 GitHub Actions: pre-commit
src/flowMC/strategy/take_steps.py
[error] 18-18: Black formatting check failed. The file was reformatted by Black to fix spacing and trailing comma issues.
[error] 112-113: Black formatting check failed. The file was reformatted to adjust line breaks and trailing commas in function calls.
[error] 124-125: Black formatting check failed. The file was reformatted to adjust line breaks and trailing commas in function calls.
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: build (3.11)
- GitHub Check: build (3.12)
🔇 Additional comments (6)
src/flowMC/resource_strategy_bundle/RQSpline_MALA.py (2)
45-45: LGTM: Parameter addition looks correct.The new
chain_batch_sizeparameter is properly added with a sensible default value of 0.
139-139: LGTM: Parameter propagation is consistent.The
chain_batch_sizeparameter is correctly passed to bothTakeSerialStepsandTakeGroupStepsconstructors, enabling the batching functionality in the underlying stepping strategies.Also applies to: 150-150
src/flowMC/resource_strategy_bundle/RQSpline_MALA_PT.py (2)
51-51: LGTM: Parameter addition is consistent.The
chain_batch_sizeparameter is added with the same default value and pattern as inRQSpline_MALA.py, maintaining consistency across the codebase.
169-169: LGTM: Parameter propagation maintains consistency.The parameter is correctly forwarded to both stepper constructors, consistent with the implementation in
RQSpline_MALA.py.Also applies to: 180-180
src/flowMC/resource/nf_model/NF_proposal.py (1)
100-102: LGTM: Correctness fix for Metropolis-Hastings acceptance logic.The variable name changes from
*_currentto*_initialcorrectly fix the acceptance ratio calculation. The acceptance decision should compare the proposed state against the initial state of the chain step, not the current state. This is a proper implementation of the Metropolis-Hastings algorithm.Also applies to: 106-107, 111-112, 114-114
src/flowMC/strategy/take_steps.py (1)
32-32: LGTM: Parameter addition is correct.The
chain_batch_sizeparameter is properly added to the constructor with a sensible default value and correctly assigned to the instance attribute.Also applies to: 42-42
1 participant
This PR introduces a new keyword
chain_batch_sizein theTakeStepsstrategy, which allows the user to batch out the chain dimension instead of processing everything at once with vmap. This reduces the memory footprint of theTakeStepsclass and allow users with smaller GPUs to run a large number of chains.The following messages are generated by copilot
This pull request introduces several updates across different files to enhance functionality and improve batch processing capabilities. The changes primarily focus on adding support for
chain_batch_sizeto manage memory constraints during computations, as well as refining variable naming for clarity in probability calculations.Batch processing enhancements:
src/flowMC/resource_strategy_bundle/RQSpline_MALA.py: Addedchain_batch_sizeas an optional parameter to the constructor and passed it to local and global step configurations. This facilitates splitting computations when processing a large number of chains. [1] [2] [3]src/flowMC/resource_strategy_bundle/RQSpline_MALA_PT.py: Similar to the above, addedchain_batch_sizeto the constructor and integrated it into local and global step configurations for parallel processing. [1] [2] [3]src/flowMC/strategy/take_steps.py: Introducedchain_batch_sizeas a class attribute and constructor parameter, and implemented logic to split computations into batches when the number of chains exceeds the batch size. This prevents memory bottlenecks during vectorized operations. [1] [2] [3] [4]Probability calculation refinements:
src/flowMC/resource/nf_model/NF_proposal.py: Renamed variables inscan_samplefromposition_currentandlog_prob_currenttoposition_initialandlog_prob_initialfor improved clarity. Adjusted related calculations to reflect these changes.Summary by CodeRabbit
New Features
Improvements
Bug Fixes