[Bug]: PyG implementation of EGNN fails with newer versions of PyG

[laserkelvin]

Expected behavior

Existing configurations should run, but changes in the PyG API appear to have deprecated some functionality.

Actual behavior

When trying to train EGNN, the following exception is thrown:

File /opt/matsciml/matsciml/models/pyg/egnn.py:177, in EGNNConv.forward(self, atom_feats, coords, edge_feats, edge_index)
    175 rel_dist = reduce(rel_coords.square(), "edges xyz -> edges ()", "sum")
    176 # combined_edge_feats = torch.cat([edge_feats, rel_dist], dim=-1)
--> 177 new_nodes, new_coords = self.propagate(
    178     edge_index,
    179     atom_feats=atom_feats,
    180     coords=coords,
    181     rel_coords=rel_coords,
    182     edge_attr=rel_dist,
    183 )
    184 return (new_nodes, new_coords)

File /opt/matsciml/matsciml/models/pyg/egnn.py:135, in EGNNConv.propagate(self, edge_index, size, **kwargs)
    133 size = self._check_input(edge_index, size)
    134 kwarg_dict = self._collect(self._user_args, edge_index, size, kwargs)
--> 135 msg_kwargs = self.inspector.distribute("message", kwarg_dict)
    136 agg_kwargs = self.inspector.distribute("aggregate", kwarg_dict)
    137 update_kwargs = self.inspector.distribute("update", kwarg_dict)

AttributeError: 'Inspector' object has no attribute 'distribute'

Seems like PyG model Inspector instances no longer have access to this distribute function. Will need to re-write the implementation to fix this.

Steps to reproduce the problem

Run either CI or the EGNN model demo script after upgrading torch_geometric==2.6.1.

Specifications

aim==3.25.0
aim-ui==3.25.0
aimrecords==0.0.7
aimrocks==0.5.2
aiofiles==24.1.0
aiohappyeyeballs==2.4.3
aiohttp==3.10.10
aioitertools==0.12.0
aiosignal==1.3.1
alabaster==0.7.16
alembic==1.13.3
anaconda-anon-usage @ file:///croot/anaconda-anon-usage_1710965072196/work
annotated-types==0.7.0
anyio==4.6.2.post1
archspec @ file:///croot/archspec_1709217642129/work
argon2-cffi==23.1.0
argon2-cffi-bindings==21.2.0
arrow==1.3.0
ase==3.23.0
asttokens @ file:///opt/conda/conda-bld/asttokens_1646925590279/work
astunparse==1.6.3
async-lru==2.0.4
attrs==24.2.0
babel==2.16.0
base58==2.0.1
bcrypt==4.2.0
beautifulsoup4 @ file:///croot/beautifulsoup4-split_1718029820055/work
bleach==6.1.0
blinker==1.8.2
boltons @ file:///work/ci_py311/boltons_1677685195580/work
boto3==1.35.44
botocore==1.35.44
Brotli @ file:///croot/brotli-split_1714483155106/work
cachetools==5.5.0
certifi==2024.8.30
cffi @ file:///croot/cffi_1714483155441/work
cfgv==3.4.0
chardet @ file:///work/ci_py311/chardet_1676830276092/work
charset-normalizer==3.4.0
click @ file:///croot/click_1698129812380/work
cloudpickle==3.1.0
colorama==0.4.6
comm==0.2.2
conda @ file:///croot/conda_1715635703388/work
conda-build @ file:///croot/conda-build_1716991285767/work
conda-content-trust @ file:///croot/conda-content-trust_1714483159009/work
conda-libmamba-solver @ file:///croot/conda-libmamba-solver_1706733287605/work/src
conda-package-handling @ file:///croot/conda-package-handling_1718138267740/work
conda_index @ file:///croot/conda-index_1719338209492/work
conda_package_streaming @ file:///croot/conda-package-streaming_1718136078615/work
ConfigArgParse==1.7
contourpy==1.3.0
cryptography @ file:///croot/cryptography_1714660666131/work
cycler==0.12.1
databricks-sdk==0.35.0
debugpy==1.8.7
decorator @ file:///opt/conda/conda-bld/decorator_1643638310831/work
defusedxml==0.7.1
Deprecated==1.2.14
dgl==2.4.0+cu121
dgllife==0.3.2
distlib==0.3.9
distro @ file:///croot/distro_1714488253808/work
dnspython==2.6.1
docker==7.1.0
docker-pycreds==0.4.0
docstring_parser==0.16
docutils==0.20.1
e3nn==0.4.4
einops==0.8.0
emmet-core==0.68.0
executing @ file:///opt/conda/conda-bld/executing_1646925071911/work
expecttest==0.2.1
fastapi==0.115.2
fastjsonschema==2.20.0
filelock @ file:///croot/filelock_1700591183607/work
Flask==3.0.3
fonttools==4.54.1
fqdn==1.5.1
frozendict @ file:///croot/frozendict_1713194832637/work
frozenlist==1.5.0
fsspec==2024.10.0
future==1.0.0
geometric-algebra-attention==0.5.1
gitdb==4.0.11
GitPython==3.1.43
gmpy2 @ file:///work/ci_py311/gmpy2_1676839849213/work
google-auth==2.35.0
graphene==3.4
graphql-core==3.2.5
graphql-relay==3.2.0
greenlet==3.1.1
gunicorn==23.0.0
h11==0.14.0
h5py==3.12.1
httpcore==1.0.6
httpx==0.27.2
hyperopt==0.2.7
hypothesis==6.108.4
identify==2.6.1
idna==3.10
ijson==3.3.0
imagesize==1.4.1
importlib_metadata==8.4.0
importlib_resources==6.4.5
iniconfig==2.0.0
ipykernel==6.29.5
ipython @ file:///croot/ipython_1718287989724/work
ipywidgets==8.1.5
isoduration==20.11.0
itsdangerous==2.2.0
jedi @ file:///croot/jedi_1721058342488/work
Jinja2==3.1.4
jmespath==1.0.1
joblib==1.4.2
json5==0.9.25
jsonargparse==4.33.2
jsonpatch @ file:///croot/jsonpatch_1714483231291/work
jsonpointer==2.1
jsonschema @ file:///croot/jsonschema_1699041609003/work
jsonschema-specifications @ file:///croot/jsonschema-specifications_1699032386549/work
jupyter==1.1.1
jupyter-console==6.6.3
jupyter-events==0.10.0
jupyter-lsp==2.2.5
jupyter_client==8.6.3
jupyter_core==5.7.2
jupyter_server==2.14.2
jupyter_server_terminals==0.5.3
jupyterlab==4.2.5
jupyterlab_pygments==0.3.0
jupyterlab_server==2.27.3
jupyterlab_widgets==3.0.13
kiwisolver==1.4.7
latexcodec==3.0.0
libarchive-c @ file:///tmp/build/80754af9/python-libarchive-c_1617780486945/work
libmambapy @ file:///croot/mamba-split_1714483352891/work/libmambapy
lightning==2.4.0
lightning-utilities==0.11.8
lintrunner==0.12.5
llvmlite==0.43.0
lmdb==1.4.1
loguru==0.7.2
mace-torch==0.3.6
maggma==0.67.0
Mako==1.3.5
Markdown==3.7
markdown-it-py==3.0.0
MarkupSafe==3.0.2
matgl==1.1.3
matplotlib==3.9.2
matplotlib-inline @ file:///work/ci_py311/matplotlib-inline_1676823841154/work
## !! Could not determine repository location
-e /opt/matsciml
matscipy==1.1.1
mdurl==0.1.2
mendeleev==0.18.1
menuinst @ file:///croot/menuinst_1718132535249/work
mistune==3.0.2
mkl-fft @ file:///croot/mkl_fft_1695058164594/work
mkl-random @ file:///croot/mkl_random_1695059800811/work
mkl-service==2.4.0
mlflow==2.17.0
mlflow-skinny==2.17.0
mongogrant==0.3.3
mongomock==4.2.0.post1
monty==2024.7.30
more-itertools @ file:///croot/more-itertools_1700662129964/work
mp-api==0.36.1
mpmath @ file:///croot/mpmath_1690848262763/work
msgpack==1.1.0
multidict==6.1.0
munch==4.0.0
nbclient==0.10.0
nbconvert==7.16.4
nbformat==5.10.4
nest-asyncio==1.6.0
networkx @ file:///croot/networkx_1720002482208/work
ninja==1.11.1.1
nodeenv==1.9.1
notebook==7.2.2
notebook_shim==0.2.4
numba==0.60.0
numpy==1.26.3
opentelemetry-api==1.27.0
opentelemetry-sdk==1.27.0
opentelemetry-semantic-conventions==0.48b0
opt-einsum-fx==0.1.4
opt_einsum==3.4.0
optree==0.12.1
orjson==3.10.9
overrides==7.7.0
packaging @ file:///croot/packaging_1720101850331/work
palettable==3.3.3
pandas==2.2.3
pandocfilters==1.5.1
paramiko==3.5.0
parso @ file:///opt/conda/conda-bld/parso_1641458642106/work
pexpect @ file:///tmp/build/80754af9/pexpect_1605563209008/work
pillow @ file:///croot/pillow_1721059439630/work
pkginfo @ file:///croot/pkginfo_1715695984887/work
platformdirs @ file:///croot/platformdirs_1692205439124/work
plotly==5.24.1
pluggy==1.5.0
pre_commit==4.0.1
prettytable==3.11.0
prometheus_client==0.21.0
prompt-toolkit @ file:///croot/prompt-toolkit_1704404351921/work
propcache==0.2.0
protobuf==5.28.2
psutil==6.1.0
ptyprocess @ file:///tmp/build/80754af9/ptyprocess_1609355006118/work/dist/ptyprocess-0.7.0-py2.py3-none-any.whl
pure-eval @ file:///opt/conda/conda-bld/pure_eval_1646925070566/work
py4j==0.10.9.7
pyarrow==17.0.0
pyasn1==0.6.1
pyasn1_modules==0.4.1
pybind11==2.13.6
pybtex==0.24.0
pycosat @ file:///croot/pycosat_1714510623388/work
pycparser @ file:///tmp/build/80754af9/pycparser_1636541352034/work
pydantic==1.10.18
pydantic-settings==2.6.0
pydantic_core==2.23.4
pydash==8.0.3
pyfiglet==0.8.post1
pyg-lib==0.4.0+pt24cu121
Pygments @ file:///croot/pygments_1684279966437/work
pymatgen==2023.7.20
pymongo==4.10.1
PyNaCl==1.5.0
pyparsing==3.2.0
PySocks @ file:///work/ci_py311/pysocks_1676822712504/work
pytest==8.3.3
pytest-dependency==0.6.0
pytest-pretty==1.2.0
python-dateutil==2.9.0.post0
python-dotenv==1.0.1
python-etcd==0.4.5
python-hostlist==1.23.0
python-json-logger==2.0.7
pytorch-lightning==2.4.0
pytz @ file:///croot/pytz_1713974312559/work
PyYAML @ file:///croot/pyyaml_1698096049011/work
pyzmq==26.2.0
rdkit==2024.3.5
referencing @ file:///croot/referencing_1699012038513/work
requests==2.32.3
RestrictedPython==7.4
rfc3339-validator==0.1.4
rfc3986-validator==0.1.1
rich==13.9.2
rowan==1.3.0.post1
rpds-py @ file:///croot/rpds-py_1698945930462/work
rsa==4.9
ruamel.yaml @ file:///work/ci_py311/ruamel.yaml_1676838772170/work
ruff==0.4.0
s3transfer==0.10.3
schema==0.7.7
scikit-learn==1.5.2
scipy==1.14.1
Send2Trash==1.8.3
sentinels==1.0.0
sentry-sdk==2.17.0
setproctitle==1.3.3
six @ file:///tmp/build/80754af9/six_1644875935023/work
smart-open==7.0.5
smmap==5.0.1
sniffio==1.3.1
snowballstemmer==2.2.0
sortedcontainers==2.4.0
soupsieve @ file:///croot/soupsieve_1696347547217/work
spglib==2.5.0
Sphinx==7.2.6
sphinx_theme_pd==0.8.3
sphinxawesome-theme==5.2.0
sphinxcontrib-applehelp==2.0.0
sphinxcontrib-devhelp==2.0.0
sphinxcontrib-htmlhelp==2.1.0
sphinxcontrib-jsmath==1.0.1
sphinxcontrib-qthelp==2.0.0
sphinxcontrib-serializinghtml==2.0.0
SQLAlchemy==2.0.36
sqlparse==0.5.1
sshtunnel==0.4.0
stack-data @ file:///opt/conda/conda-bld/stack_data_1646927590127/work
starlette==0.40.0
sympy @ file:///croot/sympy_1701397643339/work
tabulate==0.9.0
tenacity==9.0.0
tensorboardX==2.6.2.2
terminado==0.18.1
threadpoolctl==3.5.0
tinycss2==1.3.0
torch==2.4.0
torch-ema==0.3
torch-geometric==2.6.1
torch_cluster==1.6.3+pt24cu121
torch_scatter==2.1.2+pt24cu121
torch_sparse==0.6.18+pt24cu121
torch_spline_conv==1.2.2+pt24cu121
torchaudio==2.4.0
torchdata==0.7.1
torchelastic==0.2.2
torchmetrics==1.5.0
torchvision==0.19.0
tornado==6.4.1
tqdm==4.66.6
traitlets @ file:///croot/traitlets_1718227057033/work
triton==3.0.0
truststore @ file:///croot/truststore_1695244293384/work
types-dataclasses==0.6.6
types-python-dateutil==2.9.0.20241003
typeshed_client==2.7.0
typing_extensions @ file:///croot/typing_extensions_1715268824938/work
tzdata==2024.2
uncertainties==3.2.2
uri-template==1.3.0
urllib3==2.2.3
uvicorn==0.32.0
virtualenv==20.27.0
wandb==0.18.5
wcwidth @ file:///Users/ktietz/demo/mc3/conda-bld/wcwidth_1629357192024/work
webcolors==24.8.0
webencodings==0.5.1
websocket-client==1.8.0
websockets==13.1
Werkzeug==3.0.4
widgetsnbextension==4.0.13
wrapt==1.16.0
yarl==1.16.0
zipp==3.20.2
zstandard @ file:///croot/zstandard_1714677652653/work