Renamed.

Author: vbharadwaj-bk

docs/conf.py

@@ -38,7 +38,7 @@
     "torch",
     "jax",
     "openequivariance.impl_torch.extlib",
-    "openequivariance.impl_jax.extlib",
+    "openequivariance.jax.extlib",
     "openequivariance_extjax",
     "jinja2",
     "numpy",

openequivariance/openequivariance/__init__.py

@@ -72,7 +72,7 @@ def torch_ext_so_path():
 jax = None
 try:
     import openequivariance_extjax
-    import openequivariance.impl_jax as jax
+    import openequivariance.jax as jax
 except ImportError:
     pass
 

openequivariance/openequivariance/impl_jax/__init__.py

@@ -1,11 +1,11 @@
 import jax
 import numpy as np
 from functools import partial
-from openequivariance.impl_jax import extlib
+from openequivariance.jax import extlib
 from openequivariance.core.e3nn_lite import TPProblem
 from openequivariance.core.LoopUnrollTP import LoopUnrollTP
 from openequivariance.core.utils import hash_attributes
-from openequivariance.impl_jax.utils import reorder_jax
+from openequivariance.jax.utils import reorder_jax
 
 @partial(jax.custom_vjp, nondiff_argnums=(3, 4, 5))
 def forward(X, Y, W, L3_dim, irrep_dtype, attrs):

…enequivariance/impl_jax/TensorProduct.py …ce/openequivariance/jax/TensorProduct.pyopenequivariance/openequivariance/impl_jax/TensorProduct.py renamed to openequivariance/openequivariance/jax/TensorProduct.py openequivariance/openequivariance/impl_jax/TensorProduct.py renamed to openequivariance/openequivariance/jax/TensorProduct.py

@@ -1,12 +1,12 @@
 import numpy as np
 from functools import partial
 from typing import Optional
-from openequivariance.impl_jax import extlib
+from openequivariance.jax import extlib
 
 from openequivariance.core.e3nn_lite import TPProblem
 from openequivariance.core.LoopUnrollConv import LoopUnrollConv
 from openequivariance.core.utils import hash_attributes
-from openequivariance.impl_jax.utils import reorder_jax
+from openequivariance.jax.utils import reorder_jax
 
 import jax
 import jax.numpy as jnp

…uivariance/impl_jax/TensorProductConv.py …penequivariance/jax/TensorProductConv.pyopenequivariance/openequivariance/impl_jax/TensorProductConv.py renamed to openequivariance/openequivariance/jax/TensorProductConv.py openequivariance/openequivariance/impl_jax/TensorProductConv.py renamed to openequivariance/openequivariance/jax/TensorProductConv.py

@@ -0,0 +1,4 @@
+from openequivariance.jax.TensorProduct import TensorProduct as TensorProduct
+from openequivariance.jax.TensorProductConv import TensorProductConv as TensorProductConv
+
+__all__ = ["TensorProduct", "TensorProductConv"]

openequivariance/openequivariance/jax/__init__.py

@@ -47,7 +47,7 @@ def test_jax(self, request):
     def tp_and_problem(self, problem, extra_tp_constructor_args, test_jax):
         cls = oeq.TensorProduct 
         if test_jax:
-            import openequivariance.impl_jax.TensorProduct as jax_tp
+            import openequivariance.jax.TensorProduct as jax_tp
             cls = jax_tp
         tp = cls(problem, **extra_tp_constructor_args)
         return tp, problem

…equivariance/impl_jax/extlib/__init__.py …/openequivariance/jax/extlib/__init__.pyopenequivariance/openequivariance/impl_jax/extlib/__init__.py renamed to openequivariance/openequivariance/jax/extlib/__init__.py openequivariance/openequivariance/impl_jax/extlib/__init__.py renamed to openequivariance/openequivariance/jax/extlib/__init__.py

@@ -4,7 +4,6 @@
 from pytest_check import check
 
 import numpy as np
-import openequivariance
 import openequivariance as oeq
 from openequivariance.benchmark.ConvBenchmarkSuite import load_graph
 from itertools import product
@@ -60,7 +59,7 @@ def test_jax(self, request):
     def conv_object(self, request, problem, extra_conv_constructor_args, test_jax):
         cls = oeq.TensorProductConv
         if test_jax:
-            from openequivariance.impl_jax import TensorProductConv as jax_conv
+            from openequivariance.jax import TensorProductConv as jax_conv
             cls = jax_conv
 
         if request.param == "atomic":