memory efficent version of seismic - #134
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6901015e38
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if masks.is_sparse or masks.is_sparse_csr: | ||
| return self.forward_sparse(masks, return_all=return_all) |
There was a problem hiding this comment.
When callers pass a sparse CSR mask, this branch sends it into forward_sparse, but that method immediately uses COO-only operations (coalesce()/indices()). Since the condition explicitly accepts CSR masks, SparseScore(...)(csr_mask) now raises instead of computing scores; either restrict this branch to COO tensors or convert CSR masks before calling forward_sparse.
Useful? React with 👍 / 👎.
forward() dispatched CSR masks to forward_sparse(), which immediately calls coalesce()/indices() — COO-only ops — so SparseScore(...)(csr_mask) raised instead of computing scores. CSC/BSR/BSC masks took the dense path and failed there instead. forward_sparse() now converts any non-COO layout to COO up front, and forward() dispatches on `layout is not torch.strided` so every sparse layout reaches it. Also adds an explicit 2D check, since the sparse path unpacks `M, C = masks.shape`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1 participant
I modified the implementation of Seismic in my project to be more efficient in terms of memory (deleted big but useless tensors)