Fix typing for DGL 2.0

CHANGELOG.md

@@ -18,13 +18,15 @@ The rules for this file:
 
 ### Authors
 <!-- GitHub usernames of contributors to this release -->
+- lilyminium
 
 ### Reviewers
 
 ### Added
 
 ### Fixed
 <!-- Bug fixes -->
+- Fixed node and edge typing, adding DGL 2.0 compatibility (Issue #78, PR #79)
 
 ### Changed
 <!-- Changes in existing functionality -->

devtools/conda-envs/test_env.yaml

@@ -1,6 +1,7 @@
 name: openff-nagl-test
 channels:
  - openeye
+ - dglteam
  - conda-forge
  - defaults
 dependencies:

openff/nagl/molecule/_dgl/molecule.py

@@ -97,13 +97,13 @@ def from_openff(
  for offmol in offmols
  ]
  graph = dgl.batch(subgraphs)
- graph.set_batch_num_nodes(graph.batch_num_nodes().sum().reshape((-1,)))
- graph.set_batch_num_edges(
-  {
-  e_type: graph.batch_num_edges(e_type).sum().reshape((-1,))
-  for e_type in graph.canonical_etypes
- }
- )
+ n_nodes = graph.batch_num_nodes().sum().reshape((-1,))
+ graph.set_batch_num_nodes(n_nodes.type(torch.int32))
+ edges = {}
+ for e_type in graph.canonical_etypes:
+ n_edge = graph.batch_num_edges(e_type).sum().reshape((-1,))
+ edges[e_type] = n_edge.type(torch.int32)
+ graph.set_batch_num_edges(edges)
 
  mapped_smiles = offmols[0].to_smiles(mapped=True)