Added JAX reordering function.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -5,7 +5,7 @@
 from openequivariance.core.e3nn_lite import TPProblem
 from openequivariance.core.LoopUnrollTP import LoopUnrollTP
 from openequivariance.core.utils import hash_attributes
-
+from openequivariance.impl_jax.utils import reorder_jax
 
 @partial(jax.custom_vjp, nondiff_argnums=(3, 4, 5))
 def forward(X, Y, W, L3_dim, irrep_dtype, attrs):
@@ -82,7 +82,13 @@ def __call__(
         self, X: jax.numpy.ndarray, Y: jax.numpy.ndarray, W: jax.numpy.ndarray
     ) -> jax.numpy.ndarray:
         return self.forward(X, Y, W)
-    
+
+    def reorder_weights_from_e3nn(self, weights, has_batch_dim=True):
+        return reorder_jax(self.forward_schedule, weights, "forward", not self.config.shared_weights)
+
+    def reorder_weights_to_e3nn(self, weights, has_batch_dim=True):
+        return reorder_jax(self.forward_schedule, weights, "backward", not self.config.shared_weights)
+
     def forward_cpu(self, L1_in, L2_in, L3_out, weights) -> None:
         result = self.forward(
             jax.numpy.asarray(L1_in),
@@ -105,5 +111,4 @@ def backward_cpu(
         )
         L1_grad[:] = np.asarray(L1_grad_jax)
         L2_grad[:] = np.asarray(L2_grad_jax)
-        weights_grad[:] = np.asarray(weights_grad_jax)
-
+        weights_grad[:] = np.asarray(weights_grad_jax)

openequivariance/openequivariance/impl_jax/TensorProductConv.py

@@ -6,6 +6,7 @@
 from openequivariance.core.e3nn_lite import TPProblem
 from openequivariance.core.LoopUnrollConv import LoopUnrollConv
 from openequivariance.core.utils import hash_attributes
+from openequivariance.impl_jax.utils import reorder_jax
 
 import jax
 import jax.numpy as jnp
@@ -163,6 +164,12 @@ def __call__(
     ) -> jax.numpy.ndarray:
         return self.forward(X, Y, W, rows, cols, sender_perm)
 
+    def reorder_weights_from_e3nn(self, weights, has_batch_dim=True):
+        return reorder_jax(self.forward_schedule, weights, "forward", not self.config.shared_weights)
+
+    def reorder_weights_to_e3nn(self, weights, has_batch_dim=True):
+        return reorder_jax(self.forward_schedule, weights, "backward", not self.config.shared_weights)
+
     def forward_cpu(self, L1_in, L2_in, weights, L3_out, graph):
         rows = graph.rows.astype(np.int32)
         cols = graph.cols.astype(np.int32)

openequivariance/openequivariance/impl_jax/utils.py

@@ -0,0 +1,56 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+def reorder_jax_helper(schedule, weights_in, direction, has_batch_dim):
+    assert direction in ["forward", "backward"]
+
+    specs = schedule.weight_reordering_info(weights_in, has_batch_dim)
+    weights_out = jnp.zeros_like(weights_in)
+
+    for spec in specs:
+        parent_range = spec["parent_range"]
+        parent_shape = spec["parent_shape"]
+        weights_subrange = spec["weights_subrange"]
+        child_range = spec["child_range"]
+        transpose_perm = spec["transpose_perm"]
+        
+        if direction == "forward":
+            reshape_size = spec["reshape_size"]
+            
+            sliced_weights = weights_in[parent_range].reshape(parent_shape)[
+                weights_subrange
+            ]
+            
+            value_to_assign = sliced_weights.transpose(transpose_perm).reshape(reshape_size)
+            weights_out = weights_out.at[child_range].set(value_to_assign)
+
+        elif direction == "backward":
+            transpose_child_shape = spec["transpose_child_shape"]
+            child_shape = spec["child_shape"]
+
+            sliced_weights = (
+                weights_in[child_range]
+                .reshape(transpose_child_shape)
+                .transpose(transpose_perm)
+            )
+            
+            value_to_insert = sliced_weights.flatten().reshape(child_shape)
+
+            slab = weights_out[parent_range]
+            slab_reshaped = slab.reshape(parent_shape)
+            slab_reshaped = slab_reshaped.at[weights_subrange].set(value_to_insert)
+            weights_out = weights_out.at[parent_range].set(slab_reshaped.reshape(slab.shape))
+
+    return weights_out
+
+def reorder_numpy_jax_helper(schedule, weights_in, direction, has_batch_dim):
+    weights_in_jax = jnp.array(weights_in)
+    result = reorder_jax_helper(schedule, weights_in_jax, direction, has_batch_dim)
+    return np.array(result)
+
+def reorder_jax(schedule, weights_in, direction, has_batch_dim):
+    if isinstance(weights_in, (jnp.ndarray, jax.Array)):
+        return reorder_jax_helper(schedule, weights_in, direction, has_batch_dim)
+    else:
+        return reorder_numpy_jax_helper(schedule, weights_in, direction, has_batch_dim)