More plumbing.

Author: vbharadwaj-bk

openequivariance/openequivariance/core/LoopUnrollTP.py

@@ -99,6 +99,11 @@ def generate_double_backward_schedule(warps_per_block):
             "opt_level": 3,
             "irrep_dtype": dtype_to_enum[self.config.irrep_dtype],
             "weight_dtype": dtype_to_enum[self.config.weight_dtype],
+
+            # Not relevant, included for compatibility with convolution 
+            "workspace_size": 0,
+            "deterministic": 1,
+            "idx_dtype": 0
         }
 
 

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, sender_perm, workspace, L3_dim, irrep_dtype, attrs):
     forward_call = jax.ffi.ffi_call("conv_forward", 
         jax.ShapeDtypeStruct((X.shape[0], L3_dim), irrep_dtype))
@@ -53,17 +53,17 @@ def __init__(self, config: TPProblem, deterministic: bool = False, kahan: bool =
 
     def forward(
             self,
-            X: jax.ndarray, 
-            Y: jax.ndarray, 
-            W: jax.ndarray, 
-            rows: jax.ndarray, 
-            cols: jax.ndarray, 
-            sender_perm: Optional[jax.ndarray] = None) -> jax.ndarray:
+            X: jax.numpy.ndarray, 
+            Y: jax.numpy.ndarray, 
+            W: jax.numpy.ndarray, 
+            rows: jax.numpy.ndarray, 
+            cols: jax.numpy.ndarray, 
+            sender_perm: Optional[jax.numpy.ndarray] = None) -> jax.numpy.ndarray:
 
-        if self.deterministic:
+        if not self.deterministic:
             sender_perm = self.dummy_transpose_perm
         else:
-            assert sender_perm is not None, "Must provide sender_perm for non-deterministic convolutions." 
+            assert sender_perm is not None, "Must provide sender_perm for deterministic convolutions." 
 
         return forward(
             X, Y, W, 

openequivariance_extjax/src/libjax_tp_jit.cpp

@@ -144,7 +144,12 @@ std::vector<std::string> kernel_prop_keys = {
     "shared_weights", 
     "opt_level", 
     "irrep_dtype", 
-    "weight_dtype"};
+    "weight_dtype",
+
+    // Convolution only
+    "workspace_size",
+    "deterministic",
+    "idx_dtype"};
 
 std::unordered_map<string, int64_t> parse_ffi_dict(ffi::Dictionary &dict, const std::vector<string> &keys) {
     std::unordered_map<string, int64_t> result;
@@ -240,7 +245,7 @@ inline void check_tensor(const ffi::AnyBuffer &buffer,
     }
 
     if (buffer.element_type() != expected_dtype) {
-        throw std::logic_error("Datatype mismatch.");
+        throw std::logic_error("Datatype mismatch for tensor " + tensor_name);
     }
 }