Renaming + added JAX example.

vbharadwaj-bk · vbharadwaj-bk · commit 1bcea3371210 · 2025-12-24T20:44:48.000-08:00
diff --git a/docs/api.rst b/docs/api.rst
@@ -30,9 +30,9 @@ PyTorch API
     :undoc-members:
     :exclude-members: name
 
-.. autofunction:: openequivariance.torch_to_oeq_dtype
+.. autofunction:: openequivariance._torch_to_oeq_dtype
 
-.. autofunction:: openequivariance.torch_ext_so_path
+.. autofunction:: openequivariance._torch_ext_so_path
 
 JAX API
 ------------------------
diff --git a/docs/conf.py b/docs/conf.py
@@ -37,7 +37,7 @@
 autodoc_mock_imports = [
     "torch",
     "jax",
-    "openequivariance.impl_torch.extlib",
+    "openequivariance._torch.extlib",
     "openequivariance.jax.extlib",
     "openequivariance_extjax",
     "jinja2",
diff --git a/docs/supported_ops.rst b/docs/supported_ops.rst
@@ -117,7 +117,7 @@ toplevel. You can use our implementation by running
 
 .. code-block::
 
-    from openequivariance.impl_torch.symmetric_contraction import SymmetricContraction as OEQSymmetricContraction
+    from openequivariance._torch.symmetric_contraction import SymmetricContraction as OEQSymmetricContraction
 
 Some Github users report weak performance for the
 symmetric contraction backward pass; your mileage may vary.
diff --git a/openequivariance/openequivariance/__init__.py b/openequivariance/openequivariance/__init__.py
@@ -39,10 +39,11 @@ def extension_source_path():
 
 if "OEQ_NOTORCH" not in os.environ or os.environ["OEQ_NOTORCH"] != "1":
     import torch
-    from openequivariance.impl_torch.TensorProduct import TensorProduct 
-    from openequivariance.impl_torch.TensorProductConv import TensorProductConv
 
-    from openequivariance.impl_torch.extlib import torch_ext_so_path as torch_ext_so_path_internal
+    from openequivariance._torch.TensorProduct import TensorProduct 
+    from openequivariance._torch.TensorProductConv import TensorProductConv
+
+    from openequivariance._torch.extlib import torch_ext_so_path as torch_ext_so_path_internal
     from openequivariance.core.utils import torch_to_oeq_dtype
 
     torch.serialization.add_safe_globals(
diff --git a/openequivariance/openequivariance/_torch/CUEConv.py b/openequivariance/openequivariance/_torch/CUEConv.py
@@ -2,7 +2,7 @@
 import itertools
 from typing import Iterator
 
-from openequivariance.impl_torch.CUETensorProduct import CUETensorProduct
+from openequivariance._torch.CUETensorProduct import CUETensorProduct
 from openequivariance.core.ConvolutionBase import (
     ConvolutionBase,
     scatter_add_wrapper,
diff --git a/openequivariance/openequivariance/_torch/CUETensorProduct.py b/openequivariance/openequivariance/_torch/CUETensorProduct.py
diff --git a/openequivariance/openequivariance/_torch/E3NNConv.py b/openequivariance/openequivariance/_torch/E3NNConv.py
@@ -4,8 +4,8 @@
     ConvolutionBase,
     scatter_add_wrapper,
 )
-from openequivariance.impl_torch.E3NNTensorProduct import E3NNTensorProduct
-from openequivariance.impl_torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixinConv
+from openequivariance._torch.E3NNTensorProduct import E3NNTensorProduct
+from openequivariance._torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixinConv
 
 class E3NNConv(ConvolutionBase, NumpyDoubleBackwardMixinConv):
     def __init__(self, config, *, idx_dtype=np.int64, torch_op=True):
diff --git a/openequivariance/openequivariance/_torch/E3NNTensorProduct.py b/openequivariance/openequivariance/_torch/E3NNTensorProduct.py
@@ -12,7 +12,7 @@
 from openequivariance.core.TensorProductBase import TensorProductBase
 from openequivariance.core.e3nn_lite import TPProblem
 from openequivariance.benchmark.logging_utils import getLogger
-from openequivariance.impl_torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixin
+from openequivariance._torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixin
 
 TORCH_COMPILE_AUTOTUNING_DIR = pathlib.Path("triton_autotuning")
 
diff --git a/openequivariance/openequivariance/_torch/FlashTPConv.py b/openequivariance/openequivariance/_torch/FlashTPConv.py
diff --git a/openequivariance/openequivariance/_torch/NPDoubleBackwardMixin.py b/openequivariance/openequivariance/_torch/NPDoubleBackwardMixin.py
diff --git a/openequivariance/openequivariance/_torch/TensorProduct.py b/openequivariance/openequivariance/_torch/TensorProduct.py
@@ -1,12 +1,12 @@
 from openequivariance.core.LoopUnrollTP import LoopUnrollTP
 from openequivariance import TPProblem
-from openequivariance.impl_torch import extlib
+from openequivariance._torch import extlib
 import torch
 import typing
 from openequivariance.core.utils import torch_to_oeq_dtype
 from openequivariance.benchmark.logging_utils import getLogger
-from openequivariance.impl_torch.utils import reorder_torch
-from openequivariance.impl_torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixin
+from openequivariance._torch.utils import reorder_torch
+from openequivariance._torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixin
 
 logger = getLogger()
 
diff --git a/openequivariance/openequivariance/_torch/TensorProductConv.py b/openequivariance/openequivariance/_torch/TensorProductConv.py
@@ -3,8 +3,8 @@
 import numpy as np
 import torch
 
-import openequivariance.impl_torch.extlib as extlib
-from openequivariance.impl_torch.extlib import (
+import openequivariance._torch.extlib as extlib
+from openequivariance._torch.extlib import (
     JITConvImpl,
     postprocess_kernel,
     DeviceProp,
@@ -15,14 +15,14 @@
     scatter_add_wrapper,
 )
 from openequivariance.core.LoopUnrollConv import LoopUnrollConv
-from openequivariance.impl_torch.TensorProduct import TensorProduct
+from openequivariance._torch.TensorProduct import TensorProduct
 from openequivariance import TPProblem
 from openequivariance.core.utils import torch_to_oeq_dtype
 from openequivariance.core.dtype_enum import enum_to_torch_dtype
-from openequivariance.impl_torch.utils import reorder_torch
+from openequivariance._torch.utils import reorder_torch
 
 from openequivariance.benchmark.logging_utils import getLogger
-from openequivariance.impl_torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixinConv
+from openequivariance._torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixinConv
 
 
 logger = getLogger()
diff --git a/openequivariance/openequivariance/_torch/extlib/.empty b/openequivariance/openequivariance/_torch/extlib/.empty
diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py
@@ -39,9 +39,9 @@
 
 generic_module = None
 if not build_ext:
-    import openequivariance.impl_torch.extlib.generic_module
+    import openequivariance._torch.extlib.generic_module
 
-    generic_module = openequivariance.impl_torch.extlib.generic_module
+    generic_module = openequivariance._torch.extlib.generic_module
 
 elif TORCH_VERSION_CUDA_OR_HIP:
     from torch.utils.cpp_extension import library_paths, include_paths
diff --git a/openequivariance/openequivariance/_torch/symmetric_contraction/__init__.py b/openequivariance/openequivariance/_torch/symmetric_contraction/__init__.py
@@ -0,0 +1,5 @@
+from openequivariance._torch.symmetric_contraction.symmetric_contraction import (
+    SymmetricContraction,
+)
+
+__all__ = ["SymmetricContraction"]
diff --git a/openequivariance/openequivariance/_torch/symmetric_contraction/symmetric_contraction.py b/openequivariance/openequivariance/_torch/symmetric_contraction/symmetric_contraction.py
@@ -1,7 +1,7 @@
 # ruff: noqa : E402
 import torch
 
-from openequivariance.impl_torch.extlib import GroupMM_F32, GroupMM_F64
+from openequivariance._torch.extlib import GroupMM_F32, GroupMM_F64
 
 
 class GroupMM:
diff --git a/openequivariance/openequivariance/_torch/utils.py b/openequivariance/openequivariance/_torch/utils.py
diff --git a/openequivariance/openequivariance/benchmark/TestBenchmarkSuite.py b/openequivariance/openequivariance/benchmark/TestBenchmarkSuite.py
@@ -7,7 +7,7 @@
 from dataclasses import dataclass
 
 import openequivariance as oeq
-from openequivariance.impl_torch.extlib import DeviceProp
+from openequivariance._torch.extlib import DeviceProp
 from openequivariance.core.TensorProductBase import TensorProductBase
 
 from openequivariance.benchmark.logging_utils import getLogger, bcolors
diff --git a/openequivariance/openequivariance/benchmark/benchmark_utils.py b/openequivariance/openequivariance/benchmark/benchmark_utils.py
@@ -13,7 +13,7 @@
 from openequivariance.core.utils import calculate_total_nnz
 from openequivariance.core.TensorProductBase import TensorProductBase
 from openequivariance.core.e3nn_lite import TPProblem
-from openequivariance.impl_torch.CUETensorProduct import CUETensorProduct
+from openequivariance._torch.CUETensorProduct import CUETensorProduct
 from openequivariance.benchmark.logging_utils import getLogger, bcolors
 
 logger = getLogger()
diff --git a/openequivariance/openequivariance/benchmark/correctness_utils.py b/openequivariance/openequivariance/benchmark/correctness_utils.py
@@ -2,7 +2,7 @@
 
 from openequivariance.core.TensorProductBase import TensorProductBase
 from openequivariance.core.e3nn_lite import TPProblem
-from openequivariance.impl_torch.CUETensorProduct import CUETensorProduct
+from openequivariance._torch.CUETensorProduct import CUETensorProduct
 from openequivariance.benchmark.random_buffer_utils import (
     get_random_buffers_forward,
     get_random_buffers_backward,
@@ -72,7 +72,7 @@ def correctness_forward(
     prng_seed: int,
 ) -> dict:
     if reference_implementation is None:
-        from openequivariance.impl_torch.E3NNTensorProduct import E3NNTensorProduct
+        from openequivariance._torch.E3NNTensorProduct import E3NNTensorProduct
 
         reference_implementation = E3NNTensorProduct
 
@@ -116,7 +116,7 @@ def correctness_backward(
     prng_seed: int,
 ) -> dict:
     if reference_implementation is None:
-        from openequivariance.impl_torch.E3NNTensorProduct import E3NNTensorProduct
+        from openequivariance._torch.E3NNTensorProduct import E3NNTensorProduct
 
         reference_implementation = E3NNTensorProduct
 
@@ -199,7 +199,7 @@ def correctness_double_backward(
         get_random_buffers_double_backward(problem, batch_size=batch_size, prng_seed=prng_seed) 
 
     if reference_implementation is None:
-        from openequivariance.impl_torch.E3NNTensorProduct import E3NNTensorProduct
+        from openequivariance._torch.E3NNTensorProduct import E3NNTensorProduct
         reference_implementation = E3NNTensorProduct
 
     result = {"thresh": correctness_threshold, "batch_size": batch_size}
diff --git a/openequivariance/openequivariance/core/ConvolutionBase.py b/openequivariance/openequivariance/core/ConvolutionBase.py
@@ -10,7 +10,7 @@
 from openequivariance.benchmark.correctness_utils import check_similiarity
 from openequivariance.core.e3nn_lite import wigner_3j
 from openequivariance.core.utils import benchmark
-from openequivariance.impl_torch.extlib import DeviceBuffer
+from openequivariance._torch.extlib import DeviceBuffer
 
 logger = getLogger()
 
@@ -229,7 +229,7 @@ def test_correctness_forward(
         high_precision_ref=False,
     ):
         if reference_implementation is None:
-            from openequivariance.impl_torch.E3NNConv import E3NNConv
+            from openequivariance._torch.E3NNConv import E3NNConv
 
             reference_implementation = E3NNConv
 
@@ -473,7 +473,7 @@ def test_correctness_backward(
         high_precision_ref=False,
     ):
         if reference_implementation is None:
-            from openequivariance.impl_torch.E3NNConv import E3NNConv
+            from openequivariance._torch.E3NNConv import E3NNConv
 
             reference_implementation = E3NNConv
 
@@ -554,7 +554,7 @@ def test_correctness_double_backward(
             )
 
         if reference_implementation is None:
-            from openequivariance.impl_torch.E3NNConv import E3NNConv
+            from openequivariance._torch.E3NNConv import E3NNConv
             reference_implementation = E3NNConv
 
         reference_problem = self.config
diff --git a/openequivariance/openequivariance/core/TensorProductBase.py b/openequivariance/openequivariance/core/TensorProductBase.py
@@ -3,7 +3,7 @@
 from openequivariance.core.e3nn_lite import TPProblem
 from openequivariance.benchmark.logging_utils import getLogger
 from openequivariance.core.utils import benchmark
-from openequivariance.impl_torch.extlib import DeviceBuffer
+from openequivariance._torch.extlib import DeviceBuffer
 
 logger = getLogger()
 
diff --git a/openequivariance/openequivariance/core/utils.py b/openequivariance/openequivariance/core/utils.py
@@ -8,7 +8,7 @@
 import json
 import tempfile
 import hashlib
-from openequivariance.impl_torch.extlib import GPUTimer
+from openequivariance._torch.extlib import GPUTimer
 
 
 def sparse_outer_product_work(cg: np.ndarray) -> int:
diff --git a/openequivariance/openequivariance/impl_torch/symmetric_contraction/__init__.py b/openequivariance/openequivariance/impl_torch/symmetric_contraction/__init__.py
diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml
@@ -64,7 +64,7 @@ dev = [
   "cmake",
   "furo",
   "sphinx",
-  "sphinx-autobuild"
+  "sphinx-autobuild",
   "nanobind",
   "scikit-build-core"
 ]
diff --git a/tests/batch_test.py b/tests/batch_test.py
@@ -40,13 +40,13 @@ def extra_tp_constructor_args(self):
         return {}
 
     @pytest.fixture(scope="class")
-    def test_jax(self, request):
+    def with_jax(self, request):
         return request.config.getoption("--jax")
 
     @pytest.fixture(scope="class")
-    def tp_and_problem(self, problem, extra_tp_constructor_args, test_jax):
+    def tp_and_problem(self, problem, extra_tp_constructor_args, with_jax):
         cls = oeq.TensorProduct 
-        if test_jax:
+        if with_jax:
             import openequivariance.jax.TensorProduct as jax_tp
             cls = jax_tp
         tp = cls(problem, **extra_tp_constructor_args)
@@ -254,8 +254,8 @@ def problem(self, request, dtype):
 
 class TestTorchbindDisable(TestProductionModels):
     @pytest.fixture(scope="class")
-    def extra_tp_constructor_args(self, test_jax):
-        if test_jax:
+    def extra_tp_constructor_args(self, with_jax):
+        if with_jax:
             pytest.skip("N/A for JAX")
         return {"use_opaque": True}
 
@@ -270,8 +270,8 @@ def problem(self, request, dtype):
         return problem
 
     @pytest.fixture(scope="class")
-    def tp_and_problem(self, problem, extra_tp_constructor_args, test_jax):
-        if test_jax:
+    def tp_and_problem(self, problem, extra_tp_constructor_args, with_jax):
+        if with_jax:
             pytest.skip("N/A for JAX")
         else:
             tp = oeq.TensorProduct(problem, **extra_tp_constructor_args)
diff --git a/tests/benchmark.py b/tests/benchmark.py
@@ -11,14 +11,14 @@
 import numpy as np
 
 from openequivariance.benchmark.logging_utils import getLogger
-from openequivariance.impl_torch.extlib import DeviceProp
-from openequivariance.impl_torch.E3NNTensorProduct import (
+from openequivariance._torch.extlib import DeviceProp
+from openequivariance._torch.E3NNTensorProduct import (
     E3NNTensorProduct,
     E3NNTensorProductCompiledCUDAGraphs,
     E3NNTensorProductCompiledMaxAutotuneCUDAGraphs,
 )
-from openequivariance.impl_torch.TensorProduct import TensorProduct
-from openequivariance.impl_torch.CUETensorProduct import CUETensorProduct
+from openequivariance._torch.TensorProduct import TensorProduct
+from openequivariance._torch.CUETensorProduct import CUETensorProduct
 from openequivariance.benchmark.TestBenchmarkSuite import (
     TestBenchmarkSuite,
     TestDefinition,
@@ -30,15 +30,15 @@
     SingleInstruction,
 )
 
-from openequivariance.impl_torch.TensorProductConv import (
+from openequivariance._torch.TensorProductConv import (
     TensorProductConvAtomic,
     TensorProductConvDeterministic,
     TensorProductConvKahan,
     TensorProductConvScatterSum,
 )
 
-from openequivariance.impl_torch.CUEConv import CUEConv, CUEConvFused
-from openequivariance.impl_torch.FlashTPConv import FlashTPConv
+from openequivariance._torch.CUEConv import CUEConv, CUEConvFused
+from openequivariance._torch.FlashTPConv import FlashTPConv
 from openequivariance.benchmark.ConvBenchmarkSuite import ConvBenchmarkSuite, load_graph
 
 from openequivariance.benchmark.problems import (
diff --git a/tests/conv_test.py b/tests/conv_test.py
@@ -15,7 +15,6 @@
     e3tools_problems,
 )
 
-
 class ConvCorrectness:
     def thresh(self, direction):
         return {"fwd": 3e-4, "bwd": 3e-4, "double_bwd": 3e-4}[direction]
@@ -52,13 +51,13 @@ def extra_conv_constructor_args(self):
         return {}
 
     @pytest.fixture(scope="class")
-    def test_jax(self, request):
+    def with_jax(self, request):
         return request.config.getoption("--jax")
 
     @pytest.fixture(params=["atomic", "deterministic", "kahan"], scope="class")
-    def conv_object(self, request, problem, extra_conv_constructor_args, test_jax):
+    def conv_object(self, request, problem, extra_conv_constructor_args, with_jax):
         cls = oeq.TensorProductConv
-        if test_jax:
+        if with_jax:
             from openequivariance.jax import TensorProductConv as jax_conv
             cls = jax_conv
 
@@ -254,8 +253,8 @@ def conv_object(self, request, problem):
 
 class TestTorchbindDisable(TestProductionModels):
     @pytest.fixture(scope="class")
-    def extra_conv_constructor_args(self, test_jax):
-        if test_jax:
+    def extra_conv_constructor_args(self, with_jax):
+        if with_jax:
             pytest.skip("N/A for JAX")
         return {"use_opaque": True}
 
@@ -264,8 +263,8 @@ class TestTorchTo(ConvCorrectness):
     problems = [mace_problems()[0]]
 
     @pytest.fixture(params=problems, ids=lambda x: x.label, scope="class")
-    def problem(self, request, dtype, test_jax):
-        if test_jax:
+    def problem(self, request, dtype, with_jax):
+        if with_jax:
             pytest.skip("N/A for JAX")
 
         problem = request.param
diff --git a/tests/examples_test.py b/tests/examples_test.py
diff --git a/tests/export_test.py b/tests/export_test.py
