Precommit.

Author: vbharadwaj-bk

README.md

@@ -185,8 +185,7 @@ Z = tp_conv.forward(
 print(jax.numpy.linalg.norm(Z))
 
 # Test JAX JIT
-func = lambda X, Y, W, e1, e2: tp_conv.forward(X, Y, W, e1, e2) 
-jitted = jax.jit(func)
+jitted = jax.jit(lambda X, Y, W, e1, e2: tp_conv.forward(X, Y, W, e1, e2))
 print(jax.numpy.linalg.norm(jitted(X, Y, W, edge_index[0], edge_index[1])))
 ```
 

openequivariance/openequivariance/jax/TensorProduct.py

@@ -8,6 +8,7 @@
 from openequivariance.core.utils import hash_attributes
 from openequivariance.jax.utils import reorder_jax
 
+
 @partial(jax.custom_vjp, nondiff_argnums=(3, 4, 5))
 def forward(X, Y, W, L3_dim, irrep_dtype, attrs):
     forward_call = jax.ffi.ffi_call(
@@ -209,4 +210,4 @@ def double_backward_cpu(
             out_grad_jax,
         )[1]((in1_dgrad_jax, in2_dgrad_jax, weights_dgrad_jax))
 
-        return in1_grad, in2_grad, weights_grad, out_dgrad
+        return in1_grad, in2_grad, weights_grad, out_dgrad

openequivariance/openequivariance/jax/TensorProductConv.py

@@ -14,9 +14,11 @@
 
 logger = getLogger()
 
+
 def zeros_like(x):
     return jnp.zeros_like(x)
 
+
 @partial(jax.custom_vjp, nondiff_argnums=(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(
@@ -34,9 +36,7 @@ def forward_fwd(
     return out, (X, Y, W, rows, cols)
 
 
-def forward_bwd(
-    workspace, sender_perm, L3_dim, irrep_dtype, attrs, res, dZ
-):
+def forward_bwd(workspace, sender_perm, L3_dim, irrep_dtype, attrs, res, dZ):
     X, Y, W, rows, cols = res
     dX, dY, dW = backward(
         X, Y, W, dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs
@@ -60,23 +60,29 @@ def backward(X, Y, W, dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs
     return backward_call(X, Y, W, dZ, rows, cols, workspace, sender_perm, **attrs)
 
 
-def backward_fwd(
-    X, Y, W, dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs
-):
-    out = backward(
-        X, Y, W, dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs
-    )
+def backward_fwd(X, Y, W, dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs):
+    out = backward(X, Y, W, dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs)
     return out, (X, Y, W, dZ, rows, cols)
 
 
-def backward_bwd(
-    workspace, sender_perm, irrep_dtype, attrs, res, derivatives
-):
+def backward_bwd(workspace, sender_perm, irrep_dtype, attrs, res, derivatives):
     X, Y, W, dZ, rows, cols = res
     ddX, ddY, ddW = derivatives
 
     gX, gY, gW, gdZ = double_backward(
-        X, Y, W, dZ, ddX, ddY, ddW, rows, cols, workspace, sender_perm, irrep_dtype, attrs
+        X,
+        Y,
+        W,
+        dZ,
+        ddX,
+        ddY,
+        ddW,
+        rows,
+        cols,
+        workspace,
+        sender_perm,
+        irrep_dtype,
+        attrs,
     )
 
     return gX, gY, gW, gdZ, None, None
@@ -340,4 +346,4 @@ def double_backward_cpu(
             np.asarray(in2_grad),
             np.asarray(weights_grad),
             np.asarray(out_dgrad),
-        )
+        )

tests/example_test.py

@@ -162,6 +162,5 @@ def test_tutorial_jax(with_jax):
     Z = tp_conv.forward(X, Y, W, edge_index[0], edge_index[1])
     print(jax.numpy.linalg.norm(Z))
 
-    func = lambda X, Y, W, e1, e2: tp_conv.forward(X, Y, W, e1, e2) 
-    jitted = jax.jit(func)
-    print(jax.numpy.linalg.norm(jitted(X, Y, W, edge_index[0], edge_index[1])))
+    jitted = jax.jit(lambda X, Y, W, e1, e2: tp_conv.forward(X, Y, W, e1, e2))
+    print(jax.numpy.linalg.norm(jitted(X, Y, W, edge_index[0], edge_index[1])))