Skip to content

BufferList slice/negative-index support#472

Open
sudhansu-24 wants to merge 2 commits intomllam:mainfrom
sudhansu-24:fix/hierarchical-graph-viz-and-bufferlist
Open

BufferList slice/negative-index support#472
sudhansu-24 wants to merge 2 commits intomllam:mainfrom
sudhansu-24:fix/hierarchical-graph-viz-and-bufferlist

Conversation

@sudhansu-24
Copy link
Copy Markdown

@sudhansu-24 sudhansu-24 commented Mar 20, 2026

Describe your changes

Add Python sequence protocol support to BufferList and clean up downstream usage.

BufferList ergonomics / minor performance cleanup

  • Extend BufferList.getitem to support Python slicing (buffer_list[1:], buffer_list[-2:], step slices) and negative indices.

Downstream cleanup

  • Remove redundant list(...) wrapping in BaseHiGraphModel.process_step by relying on the new slice behavior.

Tests

  • Add unit tests for BufferList integer indexing, negative indexing, slicing semantics, and iteration.

Motivation / context
BufferList is 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 explicit list() conversions, matching standard Python container expectations.

The create_graph.py visualization 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

  • 🐛 Bug fix (non-breaking change that fixes an issue)
  • ✨ New feature (non-breaking change that adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

  • My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
  • I have performed a self-review of my code
  • For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
  • I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
  • I have updated the README to cover introduced code changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have given the PR a name that clearly describes the change, written in imperative form (context).
  • I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

  • the code is readable
  • the code is well tested
  • the code is documented (including return types and parameters)
  • the code is easy to maintain

Author checklist after completed review

  • I have added a line to the CHANGELOG describing this change, in a section
    reflecting type of change (add section where missing):
    • added: when you have added new functionality
    • changed: when default behaviour of the code has been changed
    • fixes: when your contribution fixes a bug
    • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

  • PR is up to date with the base branch
  • the tests pass
  • (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
  • author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
  • Once the PR is ready to be merged, squash commits and merge the PR.

- 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
@sadamov
Copy link
Copy Markdown
Collaborator

sadamov commented Mar 23, 2026

@sudhansu-24 thanks, similar to other PRs about the create_graph script I am here to warn you that we plan to soon move to weather-model-graph #83 for graph creation. That repo has beautiful (and hopefully accurate ;)) plotting in place. unless you have a strong preference to get this fixed right now, I propose to close this PR and wait for #83, let me know what you think.

@sadamov sadamov marked this pull request as draft March 23, 2026 05:32
@sudhansu-24
Copy link
Copy Markdown
Author

Hi @sadamov , thanks for the heads up!
I wasn't aware of the plan to move to weather-model-graphs (#83) that makes total sense.

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??

@sadamov
Copy link
Copy Markdown
Collaborator

sadamov commented Mar 23, 2026

@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
@sudhansu-24 sudhansu-24 changed the title Fix up-graph visualization bug and add BufferList slice support BufferList slice/negative-index support Mar 23, 2026
@sudhansu-24 sudhansu-24 marked this pull request as ready for review March 23, 2026 11:53
@sudhansu-24
Copy link
Copy Markdown
Author

@sadamov Updated #472 to focus only on BufferList slice/negative-index support and the related unit tests. removed the create_graph.py visualization changes now, marked the PR as ready for review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants