Registered the VJP rules for backward and double-backward.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_jax/TensorProductConv.py

@@ -13,7 +13,7 @@
 from openequivariance.benchmark.logging_utils import getLogger
 logger = getLogger()
 
-#@partial(jax.custom_vjp, nondiff_argnums=(3,4,5,6,7,8,9))
+@partial(jax.custom_vjp, nondiff_argnums=(3,4,5,6,7,8,9))
 def forward(X, Y, W, rows, cols, workspace, sender_perm, L3_dim, irrep_dtype, attrs):
     forward_call = jax.ffi.ffi_call("conv_forward", 
         jax.ShapeDtypeStruct((X.shape[0], L3_dim), irrep_dtype))
@@ -22,6 +22,33 @@ def forward(X, Y, W, rows, cols, workspace, sender_perm, L3_dim, irrep_dtype, at
 def forward_with_inputs(X, Y, W, rows, cols, workspace, sender_perm, L3_dim, irrep_dtype, attrs):
     return forward(X, Y, W, rows, cols, workspace, sender_perm, L3_dim, irrep_dtype, attrs), (X, Y, W, rows, cols, sender_perm, workspace)
 
+@partial(jax.custom_vjp, nondiff_argnums=(4,5,6,7,8,9))
+def backward(X, Y, W, dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs):
+    backward_call = jax.ffi.ffi_call("conv_backward", 
+        (jax.ShapeDtypeStruct(X.shape, irrep_dtype),
+         jax.ShapeDtypeStruct(Y.shape, irrep_dtype),
+         jax.ShapeDtypeStruct(W.shape, irrep_dtype)))
+    return backward_call(X, Y, W, dZ, rows, cols, workspace, sender_perm, **attrs)
+
+def backward_with_inputs(X, Y, W, dZ, rows, cols, workspace, sender_perm, L3_dim, irrep_dtype, attrs):
+    return backward(X, Y, W, dZ, rows, cols, workspace, sender_perm, L3_dim, irrep_dtype, attrs), (X, Y, W, dZ, rows, cols, sender_perm, workspace)
+
+def double_backward(rows, cols, workspace, sender_perm, irrep_dtype, attrs, inputs, derivatives):
+    double_backward_call = jax.ffi.ffi_call("conv_double_backward",
+        (
+            jax.ShapeDtypeStruct(inputs[0].shape, irrep_dtype),
+            jax.ShapeDtypeStruct(inputs[1].shape, irrep_dtype),
+            jax.ShapeDtypeStruct(inputs[2].shape, irrep_dtype),
+            jax.ShapeDtypeStruct(inputs[3].shape, irrep_dtype),
+        ))
+    return double_backward_call(*inputs, *derivatives, rows, cols, workspace, sender_perm, **attrs)
+
+def backward_autograd(rows, cols, workspace, sender_perm, L3_dim, irrep_dtype, attrs, inputs, dZ):
+    return backward(inputs[0], inputs[1], inputs[2], dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs)
+
+forward.defvjp(forward_with_inputs, backward_autograd)
+backward.defvjp(backward_with_inputs, double_backward)
+
 class TensorProductConv(LoopUnrollConv):
     def __init__(self, config: TPProblem, deterministic: bool = False, kahan: bool = False):
         dp = extlib.DeviceProp(0)
@@ -50,7 +77,6 @@ def __init__(self, config: TPProblem, deterministic: bool = False, kahan: bool =
         logger.info(f"Convolution requires {self.workspace_size // (2 ** 20)}MB of workspace.")
         self.dummy_transpose_perm = jnp.zeros((1,), dtype=jnp.int32)
 
-
     def forward(
             self,
             X: jax.numpy.ndarray,