change default cueq to ir_mul

ACEsuit · Nov 21, 2024 · a200571 · a200571
1 parent 9dca489
commit a200571
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 8 deletions.
diff --git a/mace/cli/convert_e3nn_cueq.py b/mace/cli/convert_e3nn_cueq.py
@@ -138,7 +138,7 @@ def run(
     # Add cuequivariance config
     config["cueq_config"] = CuEquivarianceConfig(
         enabled=True,
-        layout="mul_ir",
+        layout="ir_mul",
         group="O3_e3nn",
         optimize_all=True,
     )

diff --git a/mace/modules/blocks.py b/mace/modules/blocks.py
@@ -634,7 +634,6 @@ def _setup(self) -> None:
         )
 
         # Linear
-        irreps_mid = irreps_mid.simplify()
         self.irreps_out = self.target_irreps
         self.linear = Linear(
             irreps_mid,
@@ -717,7 +716,6 @@ def _setup(self) -> None:
         )
 
         # Linear
-        irreps_mid = irreps_mid.simplify()
         self.irreps_out = self.target_irreps
         self.linear = Linear(
             irreps_mid,
@@ -800,7 +798,6 @@ def _setup(self) -> None:
         )
 
         # Linear
-        irreps_mid = irreps_mid.simplify()
         self.irreps_out = self.target_irreps
         self.linear = Linear(
             irreps_mid,
@@ -898,7 +895,6 @@ def _setup(self) -> None:
         )
 
         # Linear
-        irreps_mid = irreps_mid.simplify()
         self.irreps_out = self.target_irreps
         self.linear = Linear(
             irreps_mid,
@@ -1007,7 +1003,6 @@ def _setup(self) -> None:
         )
 
         # Linear
-        irreps_mid = irreps_mid.simplify()
         self.irreps_out = self.target_irreps
         self.linear = Linear(
             irreps_mid,

diff --git a/mace/modules/wrapper_ops.py b/mace/modules/wrapper_ops.py
@@ -104,11 +104,12 @@ def __new__(
                 cue.Irreps(cueq_config.group, irreps_out),
                 layout=cueq_config.layout,
                 shared_weights=shared_weights,
+                optimize_fallback=True,
             )
             instance.original_forward = instance.forward
 
             def cuet_forward(self, x: torch.Tensor) -> torch.Tensor:
-                return self.original_forward(x, use_fallback=None)
+                return self.original_forward(x, use_fallback=True)
 
             instance.forward = types.MethodType(cuet_forward, instance)
             return instance
@@ -193,13 +194,14 @@ def __new__(
                 layout=cueq_config.layout,
                 shared_weights=shared_weights,
                 internal_weights=internal_weights,
+                optimize_fallback=True,
             )
             instance.original_forward = instance.forward
 
             def cuet_forward(
                 self, x: torch.Tensor, attrs: torch.Tensor
             ) -> torch.Tensor:
-                return self.original_forward(x, attrs, use_fallback=None)
+                return self.original_forward(x, attrs, use_fallback=True)
 
             instance.forward = types.MethodType(cuet_forward, instance)
             return instance