Forward call is working.

Author: vbharadwaj-bk

openequivariance/openequivariance/impl_jax/TensorProductConv.py

@@ -14,13 +14,13 @@
 logger = getLogger()
 
 #@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):
+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))
-    return forward_call(X, Y, W, rows, cols, sender_perm, workspace, **attrs)
+    return forward_call(X, Y, W, rows, cols, workspace, sender_perm, **attrs)
 
-def forward_with_inputs(X, Y, W, rows, cols, sender_perm, workspace, L3_dim, irrep_dtype, attrs):
-    return forward(X, Y, W, rows, cols, sender_perm, workspace, L3_dim, irrep_dtype, attrs), (X, Y, W, rows, cols, sender_perm, workspace)
+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)
 
 class TensorProductConv(LoopUnrollConv):
     def __init__(self, config: TPProblem, deterministic: bool = False, kahan: bool = False):
@@ -67,8 +67,9 @@ def forward(
 
         return forward(
             X, Y, W, 
-            rows, cols, sender_perm, 
+            rows, cols, 
             self.workspace,
+            sender_perm,
             self.L3_dim, 
             self.config.irrep_dtype, 
             self.attrs)

openequivariance_extjax/src/libjax_tp_jit.cpp

@@ -46,26 +46,20 @@ xla::ffi::DataType enum_to_xla_dtype(int64_t i){
 }
 
 inline void* data_ptr(ffi::AnyBuffer &buffer) {
-    switch (buffer.element_type()) {
-        case xla::ffi::DataType::F32:
-            return reinterpret_cast<void*>(buffer.typed_data<float>());
-        case xla::ffi::DataType::F64:
-            return reinterpret_cast<void*>(buffer.typed_data<double>());
-        case xla::ffi::DataType::S64:
-            return reinterpret_cast<void*>(buffer.typed_data<int64_t>());
-        case xla::ffi::DataType::U8:
-            return reinterpret_cast<void*>(buffer.typed_data<uint8_t>());
-        default:
-            throw logic_error("Unsupported tensor datatype!");
-    }
+    return buffer.untyped_data();
+}
+
+inline void* data_ptr(ffi::Result<ffi::AnyBuffer> &buffer) {
+    return data_ptr(*buffer);
 }
 
 inline int byte_count(ffi::AnyBuffer &buffer) {
     switch (buffer.element_type()) {
+        case xla::ffi::DataType::U32:
+        case xla::ffi::DataType::S32:
         case xla::ffi::DataType::F32:
             return 4;
         case xla::ffi::DataType::F64:
-            return 8;
         case xla::ffi::DataType::S64:
             return 8;
         case xla::ffi::DataType::U8:
@@ -75,10 +69,6 @@ inline int byte_count(ffi::AnyBuffer &buffer) {
     }
 }
 
-inline void* data_ptr(ffi::Result<ffi::AnyBuffer> &buffer) {
-    return data_ptr(*buffer);
-}
-
 #ifdef CUDA_BACKEND
 void zero_buffer(ffi::AnyBuffer &buffer) {
     cudaMemset(
@@ -245,7 +235,9 @@ inline void check_tensor(const ffi::AnyBuffer &buffer,
     }
 
     if (buffer.element_type() != expected_dtype) {
-        throw std::logic_error("Datatype mismatch for tensor " + tensor_name);
+        throw std::logic_error("Datatype mismatch for tensor " + tensor_name +
+            ". Expected datatype " + std::to_string(static_cast<int64_t>(expected_dtype)) +
+            ", got " + std::to_string(static_cast<int64_t>(buffer.element_type())));
     }
 }