misc small issues (n_trajectories -> batch_size, removed fill_value)#491
Open
josephdviviano wants to merge 3 commits intomasterfrom
Open
misc small issues (n_trajectories -> batch_size, removed fill_value)#491josephdviviano wants to merge 3 commits intomasterfrom
josephdviviano wants to merge 3 commits intomasterfrom
Conversation
The fill_value parameter in get_trajectory_pfs, get_trajectory_pbs, and get_trajectory_pfs_and_pbs was never used with a non-zero value anywhere in the codebase. Hardcode it to 0.0 and remove the parameter to simplify the API and prevent users from accidentally changing probability calculations. Also removes a dead code optimization branch (fill_value != 0.0 check in non-vectorized PF path) and the related TODO comment in get_transition_pbs. Closes #455 https://claude.ai/code/session_01NEvxmZxmjXDdSKghSeEcN8
The property n_trajectories on the Trajectories container is actually the batch size of the container. Rename it to batch_size for clarity and consistency with the rest of the codebase (e.g. RolloutContext already uses batch_size). This rename covers the property definition, all property accesses, local variables derived from it, assertions, and shape descriptions in docstrings/comments across 11 files. Tutorial CLI arguments that use n_trajectories for "total trajectories to train" are intentionally left unchanged as they represent a different concept. Closes #456 https://claude.ai/code/session_01NEvxmZxmjXDdSKghSeEcN8
Add detailed docstrings and inline comments to train_hypergrid_ddp.py and train_hypergrid_mog.py to improve readability for new users. Key additions: - Function docstrings with Args/Returns for all setup functions and main() - Inline comments explaining DDP buffer rank assignment (round-robin modulo) - Comments for WandB two-stage string broadcast pattern - On-policy detection logic explanation - Profiler schedule phase descriptions (wait/warmup/active) - Gradient all-reduce explanation in DDP sync blocks - MoG module docstring with mixture reward shaping math - f_theta classifier training and synchronization comments https://claude.ai/code/session_01NEvxmZxmjXDdSKghSeEcN8
Contributor
There was a problem hiding this comment.
Pull request overview
This PR standardizes trajectory dimension naming to batch_size across core containers/utilities and tutorial scripts, and improves the readability of the DDP/MoG example scripts through expanded docstrings/comments.
Changes:
- Replace
n_trajectoriesterminology withbatch_sizeacrossTrajectories, samplers, GFlowNet scoring, and related tutorial/test code. - Remove the
fill_valueparameter from trajectory PF/PB probability helpers and propagate the updated API through call sites/tests. - Expand documentation/comments in HyperGrid DDP and MoG training examples to better explain control flow and synchronization.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tutorials/examples/train_hypergrid_mog.py | Adds detailed module/function docs and clarifies on-policy detection + mixture reward shaping comments. |
| tutorials/examples/train_hypergrid_gafn.py | Updates shape/docs and indexing to use trajectories.batch_size. |
| tutorials/examples/train_hypergrid_ddp.py | Expands docstrings/comments describing setup, buffer-rank mapping, WandB group broadcast, profiler schedule, and grad sync. |
| tutorials/examples/train_box.py | Updates comment to reflect batch_size naming. |
| testing/test_probability_calculations.py | Updates assertions for batch_size and removes now-unsupported fill_value usage in modern calls. |
| src/gfn/utils/prob_calculations.py | Removes fill_value argument from PF/PB helpers; renames internal shape variables to batch_size. |
| src/gfn/states.py | Updates documentation to use batch_size terminology. |
| src/gfn/samplers.py | Renames n_trajectories local variables to batch_size in sampling and local-search logic. |
| src/gfn/gym/bitSequence.py | Updates naming to batch_size in generated-trajectory helper. |
| src/gfn/gflownet/sub_trajectory_balance.py | Updates docstrings and tensor shapes to use batch_size. |
| src/gfn/gflownet/base.py | Updates docs/asserts to batch_size and removes fill_value from get_pfs_and_pbs. |
| src/gfn/containers/trajectories.py | Renames public property to batch_size and updates related logic/docs. |
| src/gfn/actions.py | Updates docs to use batch_size terminology. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
| The terminating states. | ||
| """ | ||
| return self.states[self.terminating_idx - 1, torch.arange(self.n_trajectories)] | ||
| return self.states[self.terminating_idx - 1, torch.arange(self.batch_size)] |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #491 +/- ##
==========================================
+ Coverage 73.37% 73.40% +0.03%
==========================================
Files 51 51
Lines 7533 7529 -4
Branches 904 903 -1
==========================================
Hits 5527 5527
+ Misses 1674 1671 -3
+ Partials 332 331 -1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
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
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.
Description
fill_value(defaults to 0 everywhere)n_trajectoriestobatch_size