adding e3tools problems to production problems (#147)

* adding e3tools problems to production problems

* make weights external (for now)

* distribute tests appropriately into shared and non-shared

---------

Co-authored-by: Austin Glover <austin_glover@berkeley.com>

Author: asglover

openequivariance/benchmark/problems.py

@@ -149,3 +149,27 @@ def nequip_problems():
             ),
         ]
     ]
+
+
+def e3tools_problems():
+    return [
+        FCTPP(in1, in2, out, label=label, shared_weights=sw, internal_weights=iw)
+        for (in1, in2, out, label, sw, iw) in [
+            (
+                "64x0e+16x1o",
+                "1x0e+1x1o",
+                "80x0e+16x1o",
+                "e3tools_conv",
+                False,
+                False,
+            ),
+            (
+                "64x0e+16x1o",
+                "1x0e+1x1o",
+                "64x0e+16x1o",
+                "e3tools_transformer",
+                True,
+                False,  # Should be true, we don't support currently
+            ),
+        ]
+    ]

tests/conv_test.py

@@ -12,6 +12,7 @@
 from openequivariance.benchmark.problems import (
     mace_problems,
     diffdock_problems,
+    e3tools_problems,
 )
 
 
@@ -123,7 +124,9 @@ def test_tp_double_bwd(self, conv_object, graph):
 
 
 class TestProductionModels(ConvCorrectness):
-    production_model_tpps = mace_problems() + diffdock_problems()
+    production_model_tpps = (
+        mace_problems() + diffdock_problems() + [e3tools_problems()[0]]
+    )
 
     @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class")
     def problem(self, request, dtype):
@@ -219,7 +222,7 @@ def problem(self, request, dtype):
 
 
 class TestAtomicSharedWeights(ConvCorrectness):
-    problems = [mace_problems()[0], diffdock_problems()[0]]
+    problems = [mace_problems()[0], diffdock_problems()[0], e3tools_problems()[1]]
 
     def thresh(self, direction):
         return {