SymmetricContraction unusable from Python: GroupMM_F32 / GroupMM_F64 not exposed by JIT build (and AOTI .so built without bindings)

[victoriavaleeva-jpg]

Summary

from openequivariance._torch.symmetric_contraction import SymmetricContraction fails on a clean install of the released wheels (reproduced through 0.6.6). The C++ GroupMMCUDA / GroupMMHIP kernels exist, but the pybind11 bindings (GroupMM_F32, GroupMM_F64) are not registered in the file the build actually compiles, and the Python re-export in _torch/extlib/init.py is commented out. As a result SymmetricContraction cannot be constructed, and any downstream wrapper that checks for it silently falls back to a stock e3nn implementation regardless of the user's config.

Environment

openequivariance 0.6.6 (also reproduced on earlier 0.6.x)
PyTorch 2.10.0, CUDA 12.8 (also repro'd with HIP build)
Python 3.12
Linux x86_64

Root cause

PR #80 added the py::class_<GroupMM<float>> / GroupMM<double> bindings, but only in extension/kernel_wrapper.cpp. The JIT build path actually compiles extension/libtorch_tp_jit.cpp, which:

Does not #include group_mm_cuda.hpp / group_mm_hip.hpp or alias GroupMM.
Does not register GroupMM_F32 / GroupMM_F64 inside its PYBIND11_MODULE(libtorch_tp_jit, m) block.
Additionally, _torch/extlib/init.py has the two re-exports commented out:

# GroupMM_F32,
# GroupMM_F64,

so even if a build did register them, they wouldn't be importable under the public name.

The precompiled AOTI .so shipped in the wheel was also built without these bindings, so OEQ_JIT_EXTENSION=1 (forcing the JIT path) is currently the only route, and it still fails until the source is patched.

Reproduction

pip install openequivariance==0.6.6
python -c "from openequivariance._torch.symmetric_contraction import SymmetricContraction"
# ImportError: cannot import name 'GroupMM_F32' from 'openequivariance._torch.extlib'