Repository navigation
Conversation
discrete_fm.py's interpolate_edges reassigned `batch = batch[mask_i]`
(correctly slicing the per-node graph-id array down to the upper-
triangular edge positions), then immediately indexed the ALREADY-sliced
result with `mask_i` again: `self.interpolate(batch[mask_i], ...)`.
`mask_i` holds raw node indices (range [0, num_nodes)), not edge
positions, so this is indexing an edge-length array with node indices.
The sibling implementation in discrete_diffusion.py's interpolate_edges
(same method name, same docstring, same external reference) does this
correctly with a single indexing operation: `self.interpolate(batch[mask_i],
edge_attr_triu, time)`, with no premature reassignment of `batch`.
Verified the two failure modes with a standalone node/edge-index
reproduction of the slicing logic (not the full model, which needs
GPU/PyG deps unavailable in this environment):
- For graphs where num_upper_triangular_edges < num_nodes, the second
indexing raises IndexError.
- For graphs where num_upper_triangular_edges >= num_nodes (the common
case for any molecule with >3 atoms), it does not crash -- it silently
returns the wrong graph id for some edges (confirmed: an 8-node,
2-graph example mislabels one edge's graph id).
interpolate_edges is live code: module.py:250 calls it and consumes its
second return value (edge_attr_t) directly into
batch[f"{interp_param.variable_name}_t"], so every discrete edge-feature
(bond-type) flow-matching training step with more than 3 atoms in a
molecule graph hits the corrupted path, misassigning which graph's
diffusion/interpolation time each edge's corruption mask is drawn from.
Fix: remove the premature `batch = batch[mask_i]` reassignment and the
redundant second `[mask_i]` index, matching the sibling implementation
exactly.
Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com>
Author
|
Checking in — this has been open about a week with no review yet. Happy to address any feedback. |
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.
What
discrete_fm.py'sinterpolate_edgesreassignsbatch = batch[mask_i](correctly slicing the per-node graph-id array down to the upper-triangular edge positions), then immediately indexes the ALREADY-sliced result withmask_iagain:self.interpolate(batch[mask_i], ...).mask_iholds raw node indices (range[0, num_nodes)), not edge positions, so this double-indexes an edge-length array with node indices.The sibling implementation in
discrete_diffusion.py'sinterpolate_edges(same method name, same docstring, same external reference) does this correctly with a single indexing operation:self.interpolate(batch[mask_i], edge_attr_triu, time), with no premature reassignment ofbatch.interpolate_edgesis live code:module.py:250calls it and consumes its second return value (edge_attr_t) directly intobatch[f"{interp_param.variable_name}_t"], so every discrete edge-feature (bond-type) flow-matching training step with more than 3 atoms in a molecule graph hits the corrupted path, misassigning which graph's diffusion/interpolation time each edge's corruption mask is drawn from.Fix
Remove the premature
batch = batch[mask_i]reassignment and the redundant second[mask_i]index, matching the sibling implementation exactly.Verification
Reproduced the node/edge-index slicing logic standalone (not the full model, which needs GPU/PyG deps unavailable in this environment):
num_upper_triangular_edges < num_nodes, the buggy double-indexing raisesIndexError.num_upper_triangular_edges >= num_nodes(the common case for any molecule with more than 3 atoms), it doesn't crash — it silently returns the wrong graph id for some edges (confirmed with an 8-node, 2-graph example where one edge gets mislabeled into the wrong graph).