make torch geometric optional by making import of se3gnn explicit

egnn_pytorch/__init__.py

@@ -1,2 +1 @@
 from egnn_pytorch.egnn_pytorch import EGNN, EGNN_sparse, EGNN_Network
-from egnn_pytorch.se3gnn import get_sparse_adj_paths, SE3GNN_sparse

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'egnn-pytorch',
   packages = find_packages(),
-  version = '0.1.9',
+  version = '0.1.10',
   license='MIT',
   description = 'E(n)-Equivariant Graph Neural Network - Pytorch',
   author = 'Phil Wang, Eric Alcaide',

tests/test_equivariance.py

@@ -33,6 +33,16 @@ def test_egnn_equivariance():
     assert torch.allclose(coors1, (coors2 @ R + T), atol=1e-6), 'type 1 features are equivariant'
     assert not torch.allclose(feats1, feats3, atol=1e-6), 'layer must be equivariant to permutations of node order'
 
+def test_higher_dimension():
+    layer = EGNN(dim=512, edge_dim=4)
+
+    feats = torch.randn(1, 16, 512)
+    coors = torch.randn(1, 16, 5)
+    edges = torch.randn(1, 16, 16, 4)
+    mask = torch.ones(1, 16).bool()
+
+    feats, coors = layer(feats, coors, edges, mask = mask)
+    assert True
 
 def test_egnn_equivariance_with_nearest_neighbors():
     layer = EGNN(dim=512, edge_dim=1, num_nearest_neighbors=8)

tests/test_extend.py

@@ -5,7 +5,8 @@
 import torch
 
 from egnn_pytorch.utils import rot
-from egnn_pytorch import EGNN, EGNN_sparse, get_sparse_adj_paths, SE3GNN_sparse
+from egnn_pytorch import EGNN, EGNN_sparse
+from egnn_pytorch.se3gnn import get_sparse_adj_paths, SE3GNN_sparse
 
 
 def test_geom_angles():