Fix torchscipt error

Author: jwa7

python/featomic/featomic/clebsch_gordan/_dispatch.py

@@ -76,7 +76,7 @@ def permute(array, axes: List[int]):
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
-def tensordot(array_1, array_2, axes: Union[int, List[int]]):
+def tensordot(array_1, array_2, axes: List[List[int]]):
     """
     Compute tensor dot product along specified axes for arrays.
 

python/featomic/featomic/clebsch_gordan/_utils.py

@@ -365,7 +365,7 @@ def _compute_labels_full_cartesian_product(
     # Check for no shared labels dimensions
     for name in labels_1.names:
         assert name not in labels_2.names, (
-            "`labels_1` and `labels_2` must not have a" " dimension ({name}) in common"
+            "`labels_1` and `labels_2` must not have a dimension ({name}) in common"
         )
     # Create the new labels names by concatenating the names of the two input labels
     labels_names: List[str] = labels_1.names + labels_2.names

python/featomic/tests/clebsch_gordan/cartesian_spherical.py

@@ -4,7 +4,7 @@
 
 from featomic.clebsch_gordan import cartesian_to_spherical
 
-from .rotations import cartesian_rotation, WignerDReal
+from .rotations import WignerDReal, cartesian_rotation
 
 
 @pytest.fixture
@@ -382,9 +382,9 @@ def test_cartesian_to_spherical_rank_3_by_equation(cg_backend):
     assert operations.equal_metadata(l3_input, l3_reference)
     assert operations.allclose(l3_input, l3_reference)
 
+
 @pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])
 def test_cartesian_to_spherical_equivariance(cg_backend):
-
     # Define a random cartesian rank 2 tensor
     random_rank_2_arr = np.random.rand(100, 3, 3, 1)
 
@@ -396,10 +396,10 @@ def test_cartesian_to_spherical_equivariance(cg_backend):
 
     # Rotate the cartesian tensor
     random_rank_2_arr_rot = (
-            (
-            random_rank_2_arr.copy().reshape(100, 3, 3) @ R.T
-        ).transpose(0, 2, 1) @ R.T
-    ).transpose(0, 2, 1).reshape(100, 3, 3, 1)
+        ((random_rank_2_arr.copy().reshape(100, 3, 3) @ R.T).transpose(0, 2, 1) @ R.T)
+        .transpose(0, 2, 1)
+        .reshape(100, 3, 3, 1)
+    )
 
     # Build the Cartesian TMs and do cartesian to spherical
     # Build the cartesian tensor and do cartesian to spherical