remove experimental code

NVIDIA · Nov 1, 2024 · fc31ac9 · fc31ac9
1 parent e113a73
commit fc31ac9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 132 deletions.
diff --git a/cuequivariance_jax/cuequivariance_jax/primitives/tensor_product.py b/cuequivariance_jax/cuequivariance_jax/primitives/tensor_product.py
@@ -27,9 +27,6 @@
 from cuequivariance_jax.primitives.tensor_product_vanilla_impl import (
     tensor_product_vanilla_impl,
 )
-from cuequivariance_jax.primitives.tensor_product_ops_impl import (
-    tensor_product_ops_impl,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -215,12 +212,8 @@ def dispatch(
         exe: TensorProductExecution,
     ) -> list[jax.Array]:
         if platform == "cuda" and use_custom_kernels:
-            cuda_outputs = tensor_product_ops_impl(
-                *inputs, shapes=shapes, d=d, exe=exe, **options
-            )
-            if cuda_outputs is not None:
-                logger.info("using custom CUDA implementation")
-                return cuda_outputs
+            # TODO: call custom kernels here
+            pass
 
         return tensor_product_vanilla_impl(
             *inputs, shapes=shapes, d=d, exe=exe, **options

diff --git a/cuequivariance_jax/cuequivariance_jax/primitives/tensor_product_ops_impl.py b/cuequivariance_jax/cuequivariance_jax/primitives/tensor_product_ops_impl.py