All double backward tests passing.

Author: vbharadwaj-bk

openequivariance/openequivariance/benchmark/correctness_utils.py

@@ -208,11 +208,12 @@ def correctness_double_backward(
     for _, impl in enumerate([test_implementation, reference_implementation]):
         tp = instantiate_implementation(impl, problem)
         weights_reordered = tp.reorder_weights_from_e3nn(weights, has_batch_dim=not problem.shared_weights)
+        weights_dgrad_reordered = tp.reorder_weights_from_e3nn(weights_dgrad, has_batch_dim=not problem.shared_weights)
 
         if impl == CUETensorProduct and problem.shared_weights:
             weights_reordered = weights_reordered[np.newaxis, :]
 
-        in1_grad, in2_grad, weights_grad, out_dgrad = tp.double_backward_cpu(in1, in2, out_grad, weights_reordered, weights_dgrad, in1_dgrad, in2_dgrad)
+        in1_grad, in2_grad, weights_grad, out_dgrad = tp.double_backward_cpu(in1, in2, out_grad, weights_reordered, weights_dgrad_reordered, in1_dgrad, in2_dgrad)
         tensors.append(
             (   out_dgrad,
                 in1_grad,

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -127,8 +127,8 @@ def double_backward_cpu(self, in1, in2, out_grad, weights, weights_dgrad, in1_dg
         weights_dgrad_jax = jax.numpy.asarray(weights_dgrad)
 
         in1_grad, in2_grad, weights_grad, out_dgrad = jax.vjp(
-            lambda x, y, w: jax.vjp(lambda a, b, c: self.forward(a, b, c), x, y, w)[1](out_grad_jax),
-            in1_jax, in2_jax, weights_jax
+            lambda x, y, w, o: jax.vjp(lambda a, b, c: self.forward(a, b, c), x, y, w)[1](o),
+            in1_jax, in2_jax, weights_jax, out_grad_jax
         )[1]((in1_dgrad_jax, in2_dgrad_jax, weights_dgrad_jax))
 
         return in1_grad, in2_grad, weights_grad, out_dgrad