AttributeError: '_OpNamespace' 'cuequivariance' object has no attribute 'tensor_product_uniform_1d_jit'

[xiaobai249]

Dear MACE team,

When I was accelerating the model training with cuequivariance, I encountered an error: "AttributeError: '_OpNamespace' 'cuequivariance' object has no attribute 'tensor_product_uniform_1d_jit'". My MACE version is 0.3.13, pytorch version is 2.7.0, python version is 3.10.16, and cuda version is 11.8. I'm really looking forward to your reply.

The detailed error report is as follows:
Traceback (most recent call last):
File "/home/lee/anaconda3/envs/mace/bin/mace_run_train", line 8, in
sys.exit(main())
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/mace/cli/run_train.py", line 75, in main
run(args)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/mace/cli/run_train.py", line 778, in run
tools.train(
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/mace/tools/train.py", line 187, in train
valid_loss_head, eval_metrics = evaluate(
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/mace/tools/train.py", line 542, in evaluate
output = model(
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/mace/modules/models.py", line 418, in forward
node_feats, sc = interaction(
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/mace/modules/blocks.py", line 457, in forward
mji = self.conv_tp(
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/cuequivariance_torch/operations/tp_channel_wise.py", line 251, in forward
output = self.f(
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/cuequivariance_torch/primitives/segmented_polynomial.py", line 283, in forward
return self.m(inputs, input_indices, output_shapes, output_indices)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/cuequivariance_torch/primitives/segmented_polynomial_uniform_1d.py", line 254, in forward
return tensor_product_uniform_1d_jit(
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/cuequivariance_torch/primitives/segmented_polynomial_uniform_1d.py", line 62, in tensor_product_uniform_1d_jit
return torch.ops.cuequivariance.tensor_product_uniform_1d_jit(
File "/home/lee/anaconda3/envs/mace/lib/python3.10/site-packages/torch/_ops.py", line 1267, in getattr
raise AttributeError(
AttributeError: '_OpNamespace' 'cuequivariance' object has no attribute 'tensor_product_uniform_1d_jit'

[ilyes319]

Dear @xiaobai249, could you share a pip list of your environement.

[xiaobai249]

Package Version

---

anyio 4.9.0
argon2-cffi 23.1.0
argon2-cffi-bindings 21.2.0
arrow 1.3.0
ase 3.25.0
aseMolec 1.0.0
asttokens 3.0.0
async-lru 2.0.5
attrs 25.3.0
babel 2.17.0
beautifulsoup4 4.13.4
bibtexparser 1.4.3
bleach 6.2.0
blinker 1.9.0
Brotli 1.1.0
certifi 2025.10.5
cffi 1.17.1
charset-normalizer 3.4.2
click 8.2.1
colorama 0.4.6
coloredlogs 15.0.1
comm 0.2.1
ConfigArgParse 1.7
contourpy 1.3.2
cuequivariance 0.7.0
cuequivariance-ops-cu11 0.5.1
cuequivariance-ops-torch-cu11 0.5.1
cuequivariance-torch 0.7.0
cycler 0.12.1
debugpy 1.8.11
decorator 5.1.1
defusedxml 0.7.1
e3nn 0.4.4
exceptiongroup 1.2.0
executing 0.8.3
fastjsonschema 2.21.1
filelock 3.13.1
Flask 3.1.1
fonttools 4.57.0
fqdn 1.5.1
fsspec 2024.6.1
gitdb 4.0.12
GitPython 3.1.44
greenlet 3.2.3
h11 0.16.0
h2 4.2.0
h5py 3.13.0
hmmlearn 0.3.3
hpack 4.1.0
httpcore 1.0.9
httpx 0.28.1
humanfriendly 10.0
hyperframe 6.1.0
idna 3.10
importlib_metadata 8.7.0
iniconfig 2.1.0
ipykernel 6.29.5
ipython 8.30.0
ipython-genutils 0.2.0
ipywidgets 8.1.6
isoduration 20.11.0
itsdangerous 2.2.0
jedi 0.19.2
Jinja2 3.1.4
joblib 1.5.1
json5 0.12.0
jsonpointer 3.0.0
jsonschema 4.23.0
jsonschema-specifications 2025.4.1
jupyter 1.1.1
jupyter_client 8.6.3
jupyter-console 6.6.3
jupyter_contrib_core 0.4.2
jupyter_contrib_nbextensions 0.7.0
jupyter_core 5.7.2
jupyter-events 0.12.0
jupyter-highlight-selected-word 0.2.0
jupyter-lsp 2.2.5
jupyter_nbextensions_configurator 0.6.4
jupyter_server 2.15.0
jupyter_server_terminals 0.5.3
jupyterlab 4.4.1
jupyterlab_pygments 0.3.0
jupyterlab_server 2.27.3
jupyterlab_widgets 3.0.14
kiwisolver 1.4.8
lightning-utilities 0.14.3
lmdb 1.6.2
lxml 5.4.0
lz4 4.4.4
mace-torch 0.3.13
MarkupSafe 2.1.5
matplotlib 3.10.1
matplotlib-inline 0.1.6
matscipy 1.1.1
mistune 3.1.3
monty 2025.3.3
mpmath 1.3.0
munkres 1.1.4
narwhals 1.45.0
nbclient 0.10.2
nbconvert 7.16.6
nbformat 5.10.4
nest-asyncio 1.6.0
networkx 3.3
notebook 7.4.1
notebook_shim 0.2.4
numpy 1.26.4
nvidia-cublas-cu11 11.11.3.6
nvidia-cuda-cupti-cu11 11.8.87
nvidia-cuda-nvrtc-cu11 11.8.89
nvidia-cuda-runtime-cu11 11.8.89
nvidia-cudnn-cu11 9.1.0.70
nvidia-cufft-cu11 10.9.0.58
nvidia-curand-cu11 10.3.0.86
nvidia-cusolver-cu11 11.4.1.48
nvidia-cusparse-cu11 11.7.5.86
nvidia-ml-py 13.580.82
nvidia-nccl-cu11 2.21.5
nvidia-nvtx-cu11 11.8.86
olefile 0.47
opt_einsum 3.4.0
opt-einsum-fx 0.1.4
orjson 3.10.18
overrides 7.7.0
packaging 24.2
palettable 3.3.3
pandas 2.2.3
pandocfilters 1.5.1
parso 0.8.4
pexpect 4.8.0
Pillow 8.4.0
pip 25.1
platformdirs 4.3.7
plotly 6.2.0
pluggy 1.6.0
prettytable 3.16.0
prometheus_client 0.21.1
prompt-toolkit 3.0.43
psutil 5.9.0
ptyprocess 0.7.0
pure-eval 0.2.2
pycairo 1.25.0
pycparser 2.22
Pygments 2.19.1
pymatgen 2025.6.14
pynvml 13.0.1
pyparsing 3.2.3
PySocks 1.7.1
pytest 8.4.1
python-dateutil 2.9.0.post0
python-hostlist 2.2.1
python-json-logger 3.3.0
pytz 2025.2
PyYAML 6.0.2
pyzmq 26.2.0
reacnetgenerator 1.6.15
referencing 0.36.2
reportlab 3.5.68
requests 2.32.3
rfc3339-validator 0.1.4
rfc3986-validator 0.1.1
rpds-py 0.24.0
ruamel.yaml 0.18.14
ruamel.yaml.clib 0.2.12
scikit-learn 1.7.0
scipy 1.15.2
scour 0.38.2
Send2Trash 1.8.3
setuptools 78.1.1
six 1.17.0
smmap 5.0.2
sniffio 1.3.1
soupsieve 2.7
spglib 2.6.0
SQLAlchemy 2.0.41
stack-data 0.2.0
sympy 1.13.3
tabulate 0.9.0
terminado 0.18.1
threadpoolctl 3.6.0
tinycss2 1.4.0
tomli 2.2.1
torch 2.7.0+cu118
torch-dftd 0.5.1
torch-ema 0.3
torchaudio 2.7.0+cu118
torchmetrics 1.7.1
torchvision 0.22.0+cu118
tornado 6.4.2
tqdm 4.67.1
traitlets 5.14.3
triton 3.3.0
types-python-dateutil 2.9.0.20241206
typing_extensions 4.12.2
tzdata 2025.2
uncertainties 3.2.3
unicodedata2 16.0.0
uri-template 1.3.0
urllib3 2.4.0
wcwidth 0.2.13
webcolors 24.11.1
webencodings 0.5.1
websocket-client 1.8.0
Werkzeug 3.1.3
wheel 0.45.1
widgetsnbextension 4.0.14
zipp 3.23.0
zstandard 0.23.0

[ilyes319]

can you try to put all these packages to version 0.6.0?

cuequivariance 0.7.0
cuequivariance-ops-cu11 0.5.1
cuequivariance-ops-torch-cu11 0.5.1
cuequivariance-torch 0.7.0

[JiaxuanLi-THU]

Dear @xiaobai249 and @ilyes319 , I also encountered such AttributeError, and after changing the cuequivariance and cuequivariance-torch packages to version 0.6.0, this error was solved. I hope this message is helpful.