BufferList slice/negative-index support#472
Conversation
- Fix create_graph.py: plot_graph was called with pyg_down instead of pyg_up for hierarchical up-graph visualization, producing incorrect in-degree coloring - Add slice and negative-index support to BufferList.__getitem__, following the Python sequence protocol (consistent with nn.ModuleList) - Remove redundant list() conversions in BaseHiGraphModel.process_step - Add regression test for the visualization bug and unit tests for BufferList indexing
|
@sudhansu-24 thanks, similar to other PRs about the |
|
Hi @sadamov , thanks for the heads up! I'm happy to drop the create_graph.py visualization fix since wmg already handles that properly. However, the BufferList slice/negative-index support (in utils.py) & coresponding cleanup in base_hi_graph_model.py are unrelated to create_graph would it make sense to split those into a separate PR, or would you prefer to close this entirely?? |
|
@sudhansu-24 that makes sense, let's focus this one #472 on the bufferlist then |
…() wrappers - Add slice and negative-index support to BufferList.__getitem__, following the Python sequence protocol (consistent with nn.ModuleList) - Remove redundant list() conversions in BaseHiGraphModel.process_step by relying on the new slice behavior - Add unit tests for BufferList integer indexing, negative indexing, slicing semantics, and iteration
sadamov
left a comment
There was a problem hiding this comment.
No CHANGELOG entry added (author checklist item unchecked).
| if key < 0: | ||
| key += len(self) | ||
| return getattr(self, f"b{key}") |
There was a problem hiding this comment.
Bug: out-of-bounds negative index raises AttributeError instead of IndexError.
key += len(self) only runs once, so buf[-100] on a 5-element list yields key = -95 and getattr(self, "b-95") raises AttributeError. Python sequences must raise IndexError.
| if key < 0: | |
| key += len(self) | |
| return getattr(self, f"b{key}") | |
| if key < 0: | |
| key += len(self) | |
| if not (0 <= key < len(self)): | |
| raise IndexError( | |
| f"index {key} out of range for BufferList of length {len(self)}" | |
| ) | |
| return getattr(self, f"b{key}") |
| def test_negative_index(self, buffer_list): | ||
| """Negative indexing follows Python sequence convention.""" | ||
| assert torch.equal(buffer_list[-1], torch.tensor([4.0])) | ||
| assert torch.equal(buffer_list[-3], torch.tensor([2.0])) |
There was a problem hiding this comment.
Missing test for the out-of-bounds bug above.
| def test_negative_index(self, buffer_list): | |
| """Negative indexing follows Python sequence convention.""" | |
| assert torch.equal(buffer_list[-1], torch.tensor([4.0])) | |
| assert torch.equal(buffer_list[-3], torch.tensor([2.0])) | |
| def test_negative_index(self, buffer_list): | |
| """Negative indexing follows Python sequence convention.""" | |
| assert torch.equal(buffer_list[-1], torch.tensor([4.0])) | |
| assert torch.equal(buffer_list[-3], torch.tensor([2.0])) | |
| def test_out_of_bounds_raises(self, buffer_list): | |
| """Out-of-bounds access raises IndexError, not AttributeError.""" | |
| with pytest.raises(IndexError): | |
| buffer_list[5] | |
| with pytest.raises(IndexError): | |
| buffer_list[-6] |
| assert r.shape[1] == d_features | ||
|
|
||
|
|
||
| class TestBufferList: |
There was a problem hiding this comment.
TestBufferList is the only test class in the repo — all other tests use plain def test_* functions. Please flatten to match the convention.
Describe your changes
Add Python sequence protocol support to
BufferListand clean up downstream usage.BufferList ergonomics / minor performance cleanup
Downstream cleanup
Tests
Motivation / context
BufferListis used throughout the hierarchical graph models to store per-level tensors. Supporting slicing and negative indices makes hierarchical processing code more Pythonic and eliminates the need for explicitlist()conversions, matching standard Python container expectations.The
create_graph.pyvisualization fix from the original PR has been dropped per @sadamov's guidance, that code path will be superseded by the weather-model-graphs migration (#83).Dependencies
No new dependencies.
Issue Link
N/A (split from #471 — the visualization portion will be resolved by #83)
Type of change
Checklist before requesting a review
pullwith--rebaseoption if possible).Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
Author checklist after completed review
reflecting type of change (add section where missing):
Checklist for assignee