Backward test is passing.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -83,16 +83,29 @@ def __call__(
     ) -> jax.numpy.ndarray:
         return self.forward(X, Y, W)
 
-    def forward_cpu(
-        self,
-        L1_in: np.ndarray,
-        L2_in: np.ndarray,
-        L3_out: np.ndarray,
-        weights: np.ndarray,
-    ) -> None:
+    def forward_cpu(self, L1_in, L2_in, L3_out, weights) -> None:
         result = self.forward(
             jax.numpy.asarray(L1_in),
             jax.numpy.asarray(L2_in),
             jax.numpy.asarray(weights),
         )
-        L3_out[:] = np.asarray(result)
+        L3_out[:] = np.asarray(result)
+
+    def backward_cpu(
+        self, L1_in, L1_grad, L2_in, L2_grad, L3_grad, weights, weights_grad
+    ) -> None:
+        backward_fn = jax.vjp(
+            lambda X, Y, W: self.forward(X, Y, W),
+            jax.numpy.asarray(L1_in),
+            jax.numpy.asarray(L2_in),
+            jax.numpy.asarray(weights),
+        )[1]
+        L1_grad_jax, L2_grad_jax, weights_grad_jax = backward_fn(
+            jax.numpy.asarray(L3_grad)
+        )
+        L1_grad[:] = np.asarray(L1_grad_jax)
+        L2_grad[:] = np.asarray(L2_grad_jax)
+        weights_grad[:] = np.asarray(weights_grad_jax)
+
+
+