More refactoring.

Author: vbharadwaj-bk

openequivariance/openequivariance/_torch/CUETensorProduct.py

@@ -6,7 +6,7 @@
 
 from openequivariance.core.TensorProductBase import TensorProductBase
 from openequivariance.core.e3nn_lite import TPProblem
-from openequivariance.benchmark.logging_utils import getLogger
+from openequivariance.benchmark.logging import getLogger
 from openequivariance.benchmark.tpp_creation_utils import (
     ChannelwiseTPP,
     FullyConnectedTPProblem,

openequivariance/openequivariance/_torch/E3NNTensorProduct.py

@@ -11,7 +11,7 @@
 
 from openequivariance.core.TensorProductBase import TensorProductBase
 from openequivariance.core.e3nn_lite import TPProblem
-from openequivariance.benchmark.logging_utils import getLogger
+from openequivariance.benchmark.logging import getLogger
 from openequivariance._torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixin
 
 TORCH_COMPILE_AUTOTUNING_DIR = pathlib.Path("triton_autotuning")

openequivariance/openequivariance/_torch/TensorProduct.py

@@ -3,7 +3,7 @@
 from openequivariance._torch import extlib
 import torch
 from openequivariance.core.utils import torch_to_oeq_dtype, dtype_to_enum
-from openequivariance.benchmark.logging_utils import getLogger
+from openequivariance.benchmark.logging import getLogger
 from openequivariance._torch.utils import (
     reorder_torch,
     string_to_tensor,

openequivariance/openequivariance/_torch/TensorProductConv.py

@@ -23,7 +23,7 @@
     enum_to_torch_dtype,
 )
 
-from openequivariance.benchmark.logging_utils import getLogger
+from openequivariance.benchmark.logging import getLogger
 from openequivariance._torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixinConv
 
 logger = getLogger()

openequivariance/openequivariance/_torch/extlib/__init__.py

@@ -8,7 +8,7 @@
 
 import torch
 
-from openequivariance.benchmark.logging_utils import getLogger
+from openequivariance.benchmark.logging import getLogger
 
 oeq_root = str(Path(__file__).parent.parent.parent)
 

openequivariance/openequivariance/benchmark/ConvBenchmarkSuite.py

@@ -6,7 +6,12 @@
 import numpy as np
 
 import openequivariance as oeq
-from openequivariance.benchmark.logging_utils import getLogger
+from openequivariance.benchmark.correctness import (
+    correctness_backward_conv,
+    correctness_double_backward_conv,
+    correctness_forward_conv,
+)
+from openequivariance.benchmark.logging import getLogger
 from openequivariance.core.ConvolutionBase import CoordGraph
 from openequivariance.benchmark.benchmark_utils import NpEncoder
 
@@ -90,7 +95,8 @@ def run(
 
                 if direction == "forward":
                     if correctness:
-                        correctness = conv.test_correctness_forward(
+                        correctness = correctness_forward_conv(
+                            conv,
                             graph,
                             thresh=self.correctness_threshold,
                             prng_seed=self.prng_seed,
@@ -105,7 +111,8 @@ def run(
 
                 if direction == "backward":
                     if correctness:
-                        correctness = conv.test_correctness_backward(
+                        correctness = correctness_backward_conv(
+                            conv,
                             graph,
                             thresh=self.correctness_threshold,
                             prng_seed=self.prng_seed,
@@ -120,8 +127,9 @@ def run(
 
                 if direction == "double_backward":
                     if correctness:
-                        correctness = conv.test_correctness_double_backward(
-                            self.graph,
+                        correctness = correctness_double_backward_conv(
+                            conv,
+                            graph,
                             thresh=self.correctness_threshold,
                             prng_seed=self.prng_seed,
                             reference_implementation=self.reference_impl,

openequivariance/openequivariance/benchmark/TestBenchmarkSuite.py

@@ -10,9 +10,9 @@
 from openequivariance._torch.extlib import DeviceProp
 from openequivariance.core.TensorProductBase import TensorProductBase
 
-from openequivariance.benchmark.logging_utils import getLogger, bcolors
+from openequivariance.benchmark.logging import getLogger, bcolors
 from openequivariance.core.e3nn_lite import TPProblem
-from openequivariance.benchmark.correctness_utils import (
+from openequivariance.benchmark.correctness import (
     correctness_forward,
     correctness_backward,
     correctness_double_backward,

openequivariance/openequivariance/benchmark/benchmark_utils.py

@@ -15,7 +15,7 @@
 from openequivariance.core.TensorProductBase import TensorProductBase
 from openequivariance.core.e3nn_lite import TPProblem
 from openequivariance._torch.CUETensorProduct import CUETensorProduct
-from openequivariance.benchmark.logging_utils import getLogger, bcolors
+from openequivariance.benchmark.logging import getLogger, bcolors
 
 logger = getLogger()