Finished prototype of TensorProductConv.

Author: vbharadwaj-bk

openequivariance/openequivariance/core/utils.py

@@ -7,9 +7,9 @@
 
 import json
 import tempfile
+import hashlib
 from openequivariance.impl_torch.extlib import GPUTimer
 
-
 def sparse_outer_product_work(cg: np.ndarray) -> int:
     return np.sum(np.max(cg != 0, axis=2))
 
@@ -170,3 +170,13 @@ def benchmark(func, num_warmup, num_iter, mode="gpu_time", kernel_names=[]):
             time_millis[i] = kernel_time
 
     return time_millis
+
+
+def hash_attributes(attrs):
+    m = hashlib.sha256()
+
+    for key in sorted(attrs.keys()):
+        m.update(attrs[key].__repr__().encode("utf-8"))
+
+    hash = int(m.hexdigest()[:16], 16) >> 1
+    attrs["hash"] = hash

openequivariance/openequivariance/impl_jax/TensorProduct.py

@@ -7,17 +7,9 @@
 import hashlib
 from openequivariance.core.e3nn_lite import TPProblem, Irreps
 from openequivariance.core.LoopUnrollTP import LoopUnrollTP
+from openequivariance.core.utils import hash_attributes
 import jax.numpy as jnp
 
-def hash_attributes(attrs):
-    m = hashlib.sha256()
-
-    for key in sorted(attrs.keys()):
-        m.update(attrs[key].__repr__().encode("utf-8"))
-
-    hash = int(m.hexdigest()[:16], 16) >> 1
-    attrs["hash"] = hash
-
 @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("tp_forward", 
@@ -84,7 +76,6 @@ def jax_to_torch(x):
     return torch.tensor(np.asarray(x), requires_grad=True)
 
 if __name__ == "__main__":
-    tp_problem = None
     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, 

openequivariance/openequivariance/impl_jax/TensorProductConv.py

@@ -0,0 +1,54 @@
+import numpy as np
+from functools import partial
+from openequivariance.impl_jax import extlib
+
+from openequivariance.core.e3nn_lite import TPProblem, Irreps
+from openequivariance.core.LoopUnrollConv import LoopUnrollConv 
+from openequivariance.core.utils import hash_attributes
+
+import jax
+import jax.numpy as jnp
+
+from openequivariance.benchmark.logging_utils import getLogger
+logger = getLogger()
+
+class TensorProductConv(LoopUnrollConv):
+    def __init__(self, config, deterministic=False, kahan=False):
+        dp = extlib.DeviceProp(0)
+        super().__init__(
+            self,
+            config, 
+            dp, extlib.postprocess_kernel,
+            idx_dtype=np.int64,
+            torch_op=False,
+            deterministic=deterministic,
+            kahan=kahan
+        )
+
+        self.attrs = {
+            "kernel": self.jit_kernel,
+            "forward_config": vars(self.forward_schedule.launch_config),
+            "backward_config": vars(self.backward_schedule.launch_config),
+            "double_backward_config": vars(self.double_backward_schedule.launch_config),
+            "kernel_prop": self.kernelProp
+        }
+        hash_attributes(self.attrs)
+ 
+        self.weight_numel = config.weight_numel
+        self.L3_dim = self.config.irreps_out.dim
+
+        self.workspace = jnp.zeros((self.workspace_size,), dtype=jnp.uint8)
+        logger.info(f"Convolution requires {self.workspace_size // (2 ** 20)}MB of workspace.")
+        self.dummy_transpose_perm = jnp.zeros((1,), dtype=jnp.int64)
+
+
+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)
+    print("COMPLETE!")

openequivariance/openequivariance/impl_jax/extlib/__init__.py

@@ -1,6 +1,10 @@
 import jax
+import hashlib
 
 def postprocess_kernel(kernel):
+    '''
+    Only CUDA for now, so no postprocessing.
+    '''
     return kernel
 
 import openequivariance_extjax as oeq_extjax