Prepping to add tests.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -73,59 +73,4 @@ 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 jax_to_torch(x):
-    import numpy as np
-    import torch
-    return torch.tensor(np.asarray(x), requires_grad=True)
-
-if __name__ == "__main__":
-    X_ir, Y_ir, Z_ir = Irreps("1x2e"), Irreps("1x3e"), Irreps("1x2e") 
-    instructions=[(0, 0, 0, "uvu", True)]
-    problem = TPProblem(X_ir, Y_ir, Z_ir, 
-                        instructions, 
-                        shared_weights=False, 
-                        internal_weights=False)
-    tensor_product = TensorProduct(problem)
-    batch_size = 100
-
-    X = jax.random.uniform(jax.random.PRNGKey(0), (batch_size, X_ir.dim), dtype=jax.numpy.float32)
-    Y = jax.random.uniform(jax.random.PRNGKey(1), (batch_size, Y_ir.dim), dtype=jax.numpy.float32)
-    W = jax.random.uniform(jax.random.PRNGKey(2), (batch_size, tensor_product.weight_numel), dtype=jax.numpy.float32)
-    Z = tensor_product.forward(X, Y, W)
-
-    # Test forward jax vjp 
-    ctZ = jax.random.uniform(jax.random.PRNGKey(3), Z.shape, dtype=jax.numpy.float32)
-    result = jax.vjp(lambda x, y, w: tensor_product.forward(x, y, w), X, Y, W)[1](ctZ)
-
-    print("COMPLETED FORWARD PASS!")
-
-    ddX = jax.random.uniform(jax.random.PRNGKey(4), X.shape, dtype=jax.numpy.float32)
-    ddY = jax.random.uniform(jax.random.PRNGKey(5), Y.shape, dtype=jax.numpy.float32)
-    ddW = jax.random.uniform(jax.random.PRNGKey(6), W.shape, dtype=jax.numpy.float32)
-
-    result_double_backward = jax.vjp(
-        lambda x, y, w: jax.vjp(lambda a, b, c: tensor_product.forward(a, b, c), x, y, w)[1](ctZ),
-        X, Y, W
-    )[1]((ddX, ddY, ddW))
-
-    print("COMPLETED DOUBLE BACKWARD PASS!")
-
-    from e3nn import o3 
-    e3nn_tp = o3.TensorProduct(X_ir, Y_ir, Z_ir, instructions, shared_weights=False, internal_weights=False)
-    print(jax_to_torch(W).shape)
-
-    X_t = jax_to_torch(X)
-    Y_t = jax_to_torch(Y)
-    W_t = jax_to_torch(W)
-    Z_t = jax_to_torch(Z)
-    Z_e3nn = e3nn_tp(X_t, Y_t, W_t)
-    print("E3NN RESULT:", (Z_e3nn - Z_t).norm())
-
-    Z_e3nn.backward(jax_to_torch(ctZ))
-    #^^^ Print the norms of the differences in gradients instead
-    print("E3NN GRADS NORM:", (jax_to_torch(result[0]) - X_t.grad).norm(), 
-          (jax_to_torch(result[1]) - Y_t.grad).norm(), 
-          (jax_to_torch(result[2]) - W_t.grad).norm())
+        return self.forward(X, Y, W)

openequivariance/openequivariance/impl_jax/TensorProductConv.py

@@ -30,8 +30,8 @@ def backward(X, Y, W, dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs
          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 backward_with_inputs(X, Y, W, dZ, rows, cols, workspace, sender_perm, irrep_dtype, attrs):
+    return backward(X, Y, W, dZ, rows, cols, workspace, sender_perm, 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",
@@ -109,23 +109,3 @@ def __call__(self,
             sender_perm: Optional[jax.numpy.ndarray] = None) -> jax.numpy.ndarray:
         return self.forward(X, Y, W, rows, cols, sender_perm)
 
-if __name__=="__main__":
-    X_ir, Y_ir, Z_ir = Irreps("1x2e"), Irreps("1x3e"), Irreps("1x2e") 
-    instructions=[(0, 0, 0, "uvu", True)]
-    problem = TPProblem(X_ir, Y_ir, Z_ir, 
-                        instructions, 
-                        shared_weights=False, 
-                        internal_weights=False)
-
-    conv = TensorProductConv(problem, deterministic=False, kahan=False)
-
-    node_ct, nonzero_ct = 3, 4
-    X = jax.random.uniform(jax.random.PRNGKey(0), (node_ct, X_ir.dim), dtype=jax.numpy.float32)
-    Y = jax.random.uniform(jax.random.PRNGKey(1), (nonzero_ct, Y_ir.dim), dtype=jax.numpy.float32)
-    W = jax.random.uniform(jax.random.PRNGKey(2), (nonzero_ct, conv.weight_numel), dtype=jax.numpy.float32)
-    rows = jnp.array([0, 1, 1, 2], dtype=jnp.int32)
-    cols = jnp.array([1, 0, 2, 1], dtype=jnp.int32)
-    Z = conv.forward(X, Y, W, rows, cols)
-    print("Z:", Z)
-
-    print("COMPLETE!")