Skip to content

Fix/symmetric contraction groupmm bindings#198

Merged
vbharadwaj-bk merged 9 commits into
mainfrom
fix/symmetric-contraction-groupmm-bindings
May 30, 2026
Merged

Fix/symmetric contraction groupmm bindings#198
vbharadwaj-bk merged 9 commits into
mainfrom
fix/symmetric-contraction-groupmm-bindings

Conversation

@asglover
Copy link
Copy Markdown
Collaborator

Re-enable symmetric contraction / groupmm bindings and added basic pytest testing.

Copy link
Copy Markdown
Member

@vbharadwaj-bk vbharadwaj-bk left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall, looks great. Couple of things:

  1. RocM, I'm assuming hasn't been tested. Are you basically adding back in code that was taken out? In which case, no problem. Otherwise, I'm a bit wary.
  2. Have you run the test suite / installation on google collab with an NVIDIA GPU, or on a platform of your choice? I would do that, since this requires changes to the low-level extension.
  3. No problem adding cuBLAS back to the CMakeLists.txt for the stable extension, since someone wants this. I was really hoping to not do this; without it, we could package up a prebuilt wheel and avoid compilation on the user side entirely. Now, we can't do that, and worse, the user's cuBLAS should be ABI compatible with PyTorch's installed version. At some point, we should address this.
  4. Future work (not in this commit): elimination of the groupMM custom class. You can switch to a caching scheme similar to the regular OEQ kernels for the cuBLASLt state. This will enable you to compose this with torch.compile (and, incidentally, you should tell the user that this is not compatible with torch.compile (yet). Then you can add a symmetric contraction test in the export_test.py file.

Comment thread openequivariance/openequivariance/_torch/extlib/__init__.py Outdated
Comment thread openequivariance/openequivariance/_torch/extlib/__init__.py Outdated
Comment thread openequivariance/CMakeLists.txt
Comment thread openequivariance/CMakeLists.txt
Comment thread openequivariance/CMakeLists.txt
Comment thread tests/symmetric_contraction_test.py Outdated
Comment thread tests/symmetric_contraction_test.py Outdated
Comment thread tests/symmetric_contraction_test.py
@vbharadwaj-bk vbharadwaj-bk marked this pull request as ready for review May 30, 2026 21:02
@vbharadwaj-bk vbharadwaj-bk added the ci-ready Triggers CI checks for a pull request label May 30, 2026
@vbharadwaj-bk vbharadwaj-bk merged commit 152c2c2 into main May 30, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci-ready Triggers CI checks for a pull request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants