add dense in exchange block

mlff/nn/layer/so3krates_layer_sparse.py

@@ -342,11 +342,11 @@ def __call__(self, x, ev, *args, **kwargs):
         #     name='mlp_layer_1'
         # )(y)  # (N, num_features)
         # y = self.activation_fn(y)
-        # y = nn.Dense(
-        #     features=num_features + num_degrees,
-        #     kernel_init=self.last_layer_kernel_init,
-        #     name='mlp_layer_2'
-        # )(y)  # (N, num_features + num_degrees)
+        y = nn.Dense(
+            features=num_features + num_degrees,
+            kernel_init=self.last_layer_kernel_init,
+            name='mlp_layer_2'
+        )(y)  # (N, num_features + num_degrees)
         cx, cev = jnp.split(
             y,
             indices_or_sections=np.array([num_features]),

mlff/nn/representation/so3krates_sparse.py

@@ -47,7 +47,7 @@ def init_so3krates_sparse(
 
     layers = [SO3kratesLayerSparse(
         degrees=degrees,
-        use_spherical_filter=i > 0,
+        use_spherical_filter=True,
         num_heads=num_heads,
         num_features_head=num_features_head,
         qk_non_linearity=getattr(nn.activation, qk_non_linearity) if qk_non_linearity != 'identity' else lambda u: u,

tests/test_fit.py

@@ -12,7 +12,7 @@
 import pkg_resources
 
 
-def test_fit(disable_jit: bool = True):
+def test_fit(disable_jit: bool = False):
     if disable_jit:
         from jax.config import config
         config.update('jax_disable_jit', True)