From d69a80784c32645b21d1eddabcce421920b9d4a8 Mon Sep 17 00:00:00 2001
From: Guillaume Fraux <guillaume.fraux@epfl.ch>
Date: Fri, 14 Nov 2025 16:55:19 +0100
Subject: [PATCH 1/3] Update for compatibility with metatensor-operation v0.4.0

---
 python/featomic/pyproject.toml                                | 4 ++--
 .../tests/clebsch_gordan/equivariant_power_spectrum.py        | 4 +---
 .../clebsch_gordan/equivariant_power_spectrum_by_pair.py      | 4 +---
 3 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/python/featomic/pyproject.toml b/python/featomic/pyproject.toml
index c75025797..d22de5842 100644
--- a/python/featomic/pyproject.toml
+++ b/python/featomic/pyproject.toml
@@ -25,8 +25,8 @@ classifiers = [
 ]
 
 dependencies = [
-    "metatensor-core >=0.1.15,<0.2.0",
-    "metatensor-operations >=0.3.0,<0.4.0",
+    "metatensor-core >=0.1.15,<0.2",
+    "metatensor-operations >=0.4.0,<0.5",
     "wigners",
 ]
 
diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py
index 6f2ca3a77..c269202da 100644
--- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py
+++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py
@@ -163,9 +163,7 @@ def test_sample_selection() -> None:
         frame, neighbors_to_properties=True, selected_samples=label_2nd
     )
 
-    powspec_3 = metatensor.join(
-        [powspec_1, powspec_2], axis="samples", remove_tensor_name=True
-    )
+    powspec_3 = metatensor.join([powspec_1, powspec_2], axis="samples")
     powspec_4 = powspec_calc.compute(frame, neighbors_to_properties=True)
 
     assert metatensor.equal(powspec_3, powspec_4)
diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py
index dcbef8f44..cbb66a188 100644
--- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py
+++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py
@@ -152,9 +152,7 @@ def test_sample_selection() -> None:
         frame, neighbors_to_properties=True, selected_samples=label_2nd
     )
 
-    powspec_3 = metatensor.join(
-        [powspec_1, powspec_2], axis="samples", remove_tensor_name=True
-    )
+    powspec_3 = metatensor.join([powspec_1, powspec_2], axis="samples")
     powspec_4 = powspec_by_pair_calc.compute(frame, neighbors_to_properties=True)
 
     assert metatensor.equal(powspec_3, powspec_4)

From 2f44707b5034958d21b012cb6820b5898dde8446 Mon Sep 17 00:00:00 2001
From: Guillaume Fraux <guillaume.fraux@epfl.ch>
Date: Fri, 14 Nov 2025 17:13:57 +0100
Subject: [PATCH 2/3] import metatensor as mts

---
 python/featomic/tests/calculators/hypers.py   |  4 +-
 .../tests/clebsch_gordan/cg_product.py        | 86 ++++++++-----------
 .../clebsch_gordan/density_correlations.py    | 78 ++++++-----------
 .../equivariant_power_spectrum.py             | 34 ++++----
 .../equivariant_power_spectrum_by_pair.py     | 34 ++++----
 .../clebsch_gordan/density_correlations.py    |  6 +-
 6 files changed, 97 insertions(+), 145 deletions(-)

diff --git a/python/featomic/tests/calculators/hypers.py b/python/featomic/tests/calculators/hypers.py
index 1dfeed098..e973ffee3 100644
--- a/python/featomic/tests/calculators/hypers.py
+++ b/python/featomic/tests/calculators/hypers.py
@@ -1,6 +1,6 @@
 from typing import List
 
-import metatensor
+import metatensor as mts
 import pytest
 
 import featomic
@@ -128,7 +128,7 @@ def test_hypers_classes():
     )
 
     system = SystemForTests()
-    metatensor.equal_raise(with_dict.compute(system), with_classes.compute(system))
+    mts.equal_raise(with_dict.compute(system), with_classes.compute(system))
 
 
 def test_hypers_custom_classes_errors():
diff --git a/python/featomic/tests/clebsch_gordan/cg_product.py b/python/featomic/tests/clebsch_gordan/cg_product.py
index 33071ed29..d04cf0942 100644
--- a/python/featomic/tests/clebsch_gordan/cg_product.py
+++ b/python/featomic/tests/clebsch_gordan/cg_product.py
@@ -2,7 +2,7 @@
 import warnings
 from typing import List
 
-import metatensor
+import metatensor as mts
 import numpy as np
 import pytest
 from metatensor import Labels, TensorBlock, TensorMap
@@ -23,12 +23,6 @@
 except ImportError:
     HAS_TORCH = False
 
-try:
-    import metatensor.operations
-
-    HAS_METATENSOR_OPERATIONS = True
-except ImportError:
-    HAS_METATENSOR_OPERATIONS = False
 try:
     import sympy  # noqa: F401
 
@@ -103,8 +97,8 @@ def test_keys_are_matched(cg_backend):
         name for name in density.keys.names if name not in ["o3_lambda", "o3_sigma"]
     ]
     lambda_soap = calculator.compute(
-        metatensor.rename_dimension(density, "properties", "n", "n_1"),
-        metatensor.rename_dimension(density, "properties", "n", "n_2"),
+        mts.rename_dimension(density, "properties", "n", "n_1"),
+        mts.rename_dimension(density, "properties", "n", "n_2"),
         o3_lambda_1_new_name="l_1",
         o3_lambda_2_new_name="l_2",
     )
@@ -139,8 +133,8 @@ def test_key_matching_versus_manual_property_matching():
     # Compute the first lambda-SOAP with both "center_type" and "neighbor_type" in the
     # keys, so that they are matched. Then move "neighbor_type" to properties.
     lambda_soap_1 = calculator.compute(
-        metatensor.rename_dimension(density, "properties", "n", "n_1"),
-        metatensor.rename_dimension(density, "properties", "n", "n_2"),
+        mts.rename_dimension(density, "properties", "n", "n_1"),
+        mts.rename_dimension(density, "properties", "n", "n_2"),
         o3_lambda_1_new_name="l_1",
         o3_lambda_2_new_name="l_2",
         selected_keys=selected_keys,
@@ -151,16 +145,16 @@ def test_key_matching_versus_manual_property_matching():
     # "center_type" is matched.
     density = density.keys_to_properties("neighbor_type")
     lambda_soap_2 = calculator.compute(
-        metatensor.rename_dimension(
-            metatensor.rename_dimension(
+        mts.rename_dimension(
+            mts.rename_dimension(
                 density, "properties", "neighbor_type", "neighbor_1_type"
             ),
             "properties",
             "n",
             "n_1",
         ),
-        metatensor.rename_dimension(
-            metatensor.rename_dimension(
+        mts.rename_dimension(
+            mts.rename_dimension(
                 density, "properties", "neighbor_type", "neighbor_2_type"
             ),
             "properties",
@@ -174,13 +168,11 @@ def test_key_matching_versus_manual_property_matching():
 
     # Now do manual matching by slicing the properties dimension of lsoap_2,
     # i.e. identifying where "neighbor_1_type_1" == "neighbor_2_type"
-    lambda_soap_2 = metatensor.rename_dimension(
+    lambda_soap_2 = mts.rename_dimension(
         lambda_soap_2, "properties", "neighbor_1_type", "neighbor_type"
     )
 
-    lambda_soap_2 = metatensor.permute_dimensions(
-        lambda_soap_2, "properties", [0, 2, 1, 3]
-    )
+    lambda_soap_2 = mts.permute_dimensions(lambda_soap_2, "properties", [0, 2, 1, 3])
     new_blocks = []
     for block in lambda_soap_2:
         properties_filter = block.properties.column(
@@ -202,9 +194,7 @@ def test_key_matching_versus_manual_property_matching():
     lambda_soap_2 = TensorMap(lambda_soap_2.keys, new_blocks)
 
     # Check for equivalence. Sorting of metadata required here.
-    assert metatensor.allclose(
-        metatensor.sort(lambda_soap_1), metatensor.sort(lambda_soap_2)
-    )
+    assert mts.allclose(mts.sort(lambda_soap_1), mts.sort(lambda_soap_2))
 
 
 def test_error_same_name_properties():
@@ -243,8 +233,8 @@ def test_sufficient_max_angular_with_angular_cutoff():
     density = spherical_expansion(frames)
 
     calculator.compute(  # no error
-        metatensor.rename_dimension(density, "properties", "n", "n_1"),
-        metatensor.rename_dimension(density, "properties", "n", "n_2"),
+        mts.rename_dimension(density, "properties", "n", "n_1"),
+        mts.rename_dimension(density, "properties", "n", "n_2"),
         o3_lambda_1_new_name="l_1",
         o3_lambda_2_new_name="l_2",
         selected_keys=Labels(
@@ -271,8 +261,8 @@ def test_error_insufficient_max_angular():
     )
     with pytest.raises(ValueError, match=error_msg):
         calculator.compute(
-            metatensor.rename_dimension(density, "properties", "n", "n_1"),
-            metatensor.rename_dimension(density, "properties", "n", "n_2"),
+            mts.rename_dimension(density, "properties", "n", "n_1"),
+            mts.rename_dimension(density, "properties", "n", "n_2"),
             o3_lambda_1_new_name="l_1",
             o3_lambda_2_new_name="l_2",
             selected_keys=Labels(
@@ -295,12 +285,12 @@ def test_same_samples():
     # Compute the density
     density = spherical_expansion(frames)
     density = density.keys_to_properties("neighbor_type")
-    density_1 = metatensor.rename_dimension(density, "properties", "n", "n_1")
-    density_1 = metatensor.rename_dimension(
+    density_1 = mts.rename_dimension(density, "properties", "n", "n_1")
+    density_1 = mts.rename_dimension(
         density_1, "properties", "neighbor_type", "neighbor_1_type"
     )
-    density_2 = metatensor.rename_dimension(density, "properties", "n", "n_2")
-    density_2 = metatensor.rename_dimension(
+    density_2 = mts.rename_dimension(density, "properties", "n", "n_2")
+    density_2 = mts.rename_dimension(
         density_2, "properties", "neighbor_type", "neighbor_2_type"
     )
 
@@ -336,24 +326,20 @@ def test_different_samples():
 
     # Compute the density
     density = spherical_expansion(frames)
-    density = metatensor.rename_dimension(
-        density, "keys", "center_type", "first_atom_type"
-    )
-    density = metatensor.rename_dimension(
-        density, "keys", "neighbor_type", "second_atom_type"
-    )
-    density = metatensor.rename_dimension(density, "samples", "atom", "first_atom")
+    density = mts.rename_dimension(density, "keys", "center_type", "first_atom_type")
+    density = mts.rename_dimension(density, "keys", "neighbor_type", "second_atom_type")
+    density = mts.rename_dimension(density, "samples", "atom", "first_atom")
     density = density.keys_to_properties("second_atom_type")
-    density_2 = metatensor.rename_dimension(density, "properties", "n", "n_2")
-    density_2 = metatensor.rename_dimension(
+    density_2 = mts.rename_dimension(density, "properties", "n", "n_2")
+    density_2 = mts.rename_dimension(
         density_2, "properties", "second_atom_type", "second_atom_2_type"
     )
 
     # Compute pair density
     pair_density = spherical_expansion_by_pair(frames)
     pair_density = pair_density.keys_to_properties("second_atom_type")
-    density_1 = metatensor.rename_dimension(pair_density, "properties", "n", "n_1")
-    density_1 = metatensor.rename_dimension(
+    density_1 = mts.rename_dimension(pair_density, "properties", "n", "n_1")
+    density_1 = mts.rename_dimension(
         pair_density, "properties", "second_atom_type", "second_atom_1_type"
     )
 
@@ -367,9 +353,9 @@ def test_different_samples():
     )
 
     # Move everything but "o3_lambda" and "center_type" to properties
-    density_1 = metatensor.sort(density_1)
-    lsoap = metatensor.sort(lsoap.keys_to_properties(new_o3_lambda_names))
-    lsoap = metatensor.sum_over_samples(
+    density_1 = mts.sort(density_1)
+    lsoap = mts.sort(lsoap.keys_to_properties(new_o3_lambda_names))
+    lsoap = mts.sum_over_samples(
         lsoap, ["second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"]
     )
 
@@ -429,8 +415,8 @@ def test_device_dtype(dtype, device):
 
     # just checking that the code runs without error
     calculator.compute(
-        metatensor.rename_dimension(density, "properties", "n", "n_1"),
-        metatensor.rename_dimension(density, "properties", "n", "n_2"),
+        mts.rename_dimension(density, "properties", "n", "n_1"),
+        mts.rename_dimension(density, "properties", "n", "n_2"),
         o3_lambda_1_new_name="l_1",
         o3_lambda_2_new_name="l_2",
     )
@@ -453,8 +439,8 @@ def test_dense_sparse_agree():
         )
         results.append(
             calculator.compute(
-                metatensor.rename_dimension(density, "properties", "n", "n_1"),
-                metatensor.rename_dimension(density, "properties", "n", "n_2"),
+                mts.rename_dimension(density, "properties", "n", "n_1"),
+                mts.rename_dimension(density, "properties", "n", "n_2"),
                 o3_lambda_1_new_name="l_1",
                 o3_lambda_2_new_name="l_2",
                 selected_keys=Labels(
@@ -464,5 +450,5 @@ def test_dense_sparse_agree():
             )
         )
 
-    assert metatensor.equal_metadata(results[0], results[1])
-    assert metatensor.allclose(results[0], results[1])
+    assert mts.equal_metadata(results[0], results[1])
+    assert mts.allclose(results[0], results[1])
diff --git a/python/featomic/tests/clebsch_gordan/density_correlations.py b/python/featomic/tests/clebsch_gordan/density_correlations.py
index 99b9c9dc8..647bf2359 100644
--- a/python/featomic/tests/clebsch_gordan/density_correlations.py
+++ b/python/featomic/tests/clebsch_gordan/density_correlations.py
@@ -1,6 +1,6 @@
 from typing import List
 
-import metatensor
+import metatensor as mts
 import numpy as np
 import pytest
 from metatensor import Labels, TensorBlock, TensorMap
@@ -20,12 +20,6 @@
 from ase import io  # noqa: E402, F401
 
 
-try:
-    import metatensor.operations
-
-    HAS_METATENSOR_OPERATIONS = True
-except ImportError:
-    HAS_METATENSOR_OPERATIONS = False
 try:
     import sympy  # noqa: F401
 
@@ -143,9 +137,7 @@ def get_norm(tensor: TensorMap):
     as a sum over lambda, sigma, and m.
     """
     # Check that there is only one sample
-    samples = metatensor.unique_metadata(
-        tensor, "samples", ["system", "atom", "center_type"]
-    )
+    samples = mts.unique_metadata(tensor, "samples", ["system", "atom", "center_type"])
     assert len(samples) == 1
 
     norm = 0.0
@@ -164,10 +156,7 @@ def get_norm(tensor: TensorMap):
 # ============ Test equivariance ============
 
 
-@pytest.mark.skipif(
-    not HAS_SYMPY or not HAS_METATENSOR_OPERATIONS,
-    reason="SymPy or metatensor-operations are not installed",
-)
+@pytest.mark.skipif(not HAS_SYMPY, reason="SymPy is not installed")
 def test_so3_equivariance():
     """
     Tests that the output of :py:func:`correlate_density` is equivariant under
@@ -195,13 +184,10 @@ def test_so3_equivariance():
     nu_3_so3 = calculator.compute(density_so3)
 
     nu_3_transf = wig.transform_tensormap_so3(nu_3)
-    assert metatensor.allclose(nu_3_transf, nu_3_so3)
+    assert mts.allclose(nu_3_transf, nu_3_so3)
 
 
-@pytest.mark.skipif(
-    not HAS_SYMPY or not HAS_METATENSOR_OPERATIONS,
-    reason="SymPy or metatensor-operations are not installed",
-)
+@pytest.mark.skipif(not HAS_SYMPY, reason="SymPy is not installed")
 def test_o3_equivariance():
     """
     Tests that the output of :py:func:`correlate_density` is equivariant under
@@ -229,15 +215,12 @@ def test_o3_equivariance():
     nu_3_o3 = calculator.compute(density_o3, selected_keys=selected_keys)
 
     nu_3_transf = wig.transform_tensormap_o3(nu_3)
-    assert metatensor.allclose(nu_3_transf, nu_3_o3)
+    assert mts.allclose(nu_3_transf, nu_3_o3)
 
 
 # ============ Test lambda-SOAP vs PowerSpectrum ============
 
 
-@pytest.mark.skipif(
-    not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed"
-)
 def test_lambda_soap_vs_powerspectrum():
     """
     Tests for exact equivalence between the invariant block of a generated
@@ -285,15 +268,12 @@ def test_lambda_soap_vs_powerspectrum():
         )
     lambda_soap = TensorMap(keys=keys, blocks=blocks)
 
-    assert metatensor.allclose(lambda_soap, ps)
+    assert mts.allclose(lambda_soap, ps)
 
 
 # ============ Test norm preservation  ============
 
 
-@pytest.mark.skipif(
-    not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed"
-)
 def test_correlate_density_norm():
     """
     Checks \\|\\rho^\\nu\\| = \\|\\rho\\|^\\nu in the case where l lists are not
@@ -330,7 +310,7 @@ def test_correlate_density_norm():
 
     # The norm should be calculated for each sample. First find the unique
     # samples
-    unique_samples = metatensor.unique_metadata(
+    unique_samples = mts.unique_metadata(
         ps, "samples", names=["system", "atom", "center_type"]
     )
     selections = [
@@ -344,9 +324,9 @@ def test_correlate_density_norm():
     norm_ps_sorted = 0.0
     for selection in selections:
         # Slice the TensorMaps
-        nu1_sliced = metatensor.slice(density, "samples", selection=selection)
-        ps_sliced = metatensor.slice(ps, "samples", selection=selection)
-        ps_sorted_sliced = metatensor.slice(ps_sorted, "samples", selection=selection)
+        nu1_sliced = mts.slice(density, "samples", selection=selection)
+        ps_sliced = mts.slice(ps, "samples", selection=selection)
+        ps_sorted_sliced = mts.slice(ps_sorted, "samples", selection=selection)
 
         # Calculate norms
         norm_nu1 += get_norm(nu1_sliced) ** (n_correlations + 1)
@@ -452,9 +432,6 @@ def test_clebsch_gordan_orthogonality(l1, l2, arrays_backend):
     )
 
 
-@pytest.mark.skipif(
-    not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed"
-)
 def test_correlate_density_dense_sparse_agree():
     """
     Tests for agreement between nu=2 tensors built using both sparse and dense
@@ -479,15 +456,12 @@ def test_correlate_density_dense_sparse_agree():
     n_body_sparse = calculator_sparse.compute(density)
     n_body_dense = calculator_dense.compute(density)
 
-    assert metatensor.allclose(n_body_sparse, n_body_dense, atol=1e-8, rtol=1e-8)
+    assert mts.allclose(n_body_sparse, n_body_dense, atol=1e-8, rtol=1e-8)
 
 
 # ============ Test metadata  ============
 
 
-@pytest.mark.skipif(
-    not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed"
-)
 def test_correlate_density_metadata_agree():
     """
     Tests that the metadata of outputs from :py:func:`correlate_density` and
@@ -511,7 +485,7 @@ def test_correlate_density_metadata_agree():
         # Build higher body order tensor without CG computation - i.e. metadata
         # only
         nu_x_metadata_only = calculator.compute_metadata(nu_1)
-        assert metatensor.equal_metadata(nu_x, nu_x_metadata_only)
+        assert mts.equal_metadata(nu_x, nu_x_metadata_only)
 
 
 @pytest.mark.parametrize(
@@ -583,29 +557,25 @@ def test_summed_powerspectrum_by_pair_equals_powerspectrum():
 
     # Generate density and rename dimensions ready for correlation
     density = spherical_expansion(frames)
-    density = metatensor.rename_dimension(
-        density, "keys", "center_type", "first_atom_type"
-    )
-    density = metatensor.rename_dimension(
-        density, "keys", "neighbor_type", "second_atom_type"
-    )
-    density = metatensor.rename_dimension(density, "samples", "atom", "first_atom")
+    density = mts.rename_dimension(density, "keys", "center_type", "first_atom_type")
+    density = mts.rename_dimension(density, "keys", "neighbor_type", "second_atom_type")
+    density = mts.rename_dimension(density, "samples", "atom", "first_atom")
     density = density.keys_to_properties("second_atom_type")
 
     # Calculate power spectrum with DensityCorrelations
     power_spectrum = density_correlations(density)
 
     # Rename dimensions ready for manual CG tensor product by ClebschGordanProduct
-    density = metatensor.rename_dimension(density, "properties", "n", "n_1")
-    density = metatensor.rename_dimension(
+    density = mts.rename_dimension(density, "properties", "n", "n_1")
+    density = mts.rename_dimension(
         density, "properties", "second_atom_type", "second_atom_1_type"
     )
 
     # Generate pair density
     pair_density = spherical_expansion_by_pair(frames)
     pair_density = pair_density.keys_to_properties("second_atom_type")
-    pair_density = metatensor.rename_dimension(pair_density, "properties", "n", "n_2")
-    pair_density = metatensor.rename_dimension(
+    pair_density = mts.rename_dimension(pair_density, "properties", "n", "n_2")
+    pair_density = mts.rename_dimension(
         pair_density, "properties", "second_atom_type", "second_atom_2_type"
     )
 
@@ -618,15 +588,15 @@ def test_summed_powerspectrum_by_pair_equals_powerspectrum():
     )
 
     # Sum, sort, check equivalence
-    power_spec_by_pair_reduced = metatensor.sum_over_samples(
+    power_spec_by_pair_reduced = mts.sum_over_samples(
         power_spectrum_by_pair,
         ["second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"],
     )
 
     # Compare
-    metatensor.allclose_raise(
-        metatensor.sort(power_spectrum),
-        metatensor.sort(power_spec_by_pair_reduced),
+    mts.allclose_raise(
+        mts.sort(power_spectrum),
+        mts.sort(power_spec_by_pair_reduced),
     )
 
 
diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py
index c269202da..d73fa55b0 100644
--- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py
+++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py
@@ -1,6 +1,6 @@
 from typing import List
 
-import metatensor
+import metatensor as mts
 import numpy as np
 import pytest
 from metatensor import Labels, TensorBlock, TensorMap
@@ -96,13 +96,13 @@ def test_equivariant_power_spectrum_vs_powerspectrum():
     ps_2 = TensorMap(keys=keys, blocks=blocks)
 
     # Permute properties dimension to match ps_1 and sort
-    ps_2 = metatensor.sort(
-        metatensor.permute_dimensions(ps_2, "properties", [2, 0, 3, 1, 4]),
+    ps_2 = mts.sort(
+        mts.permute_dimensions(ps_2, "properties", [2, 0, 3, 1, 4]),
         "properties",
     )
 
-    metatensor.equal_metadata_raise(ps_1, ps_2)
-    metatensor.allclose_raise(ps_1, ps_2)
+    mts.equal_metadata_raise(ps_1, ps_2)
+    mts.allclose_raise(ps_1, ps_2)
 
 
 def test_equivariant_power_spectrum_neighbors_to_properties():
@@ -127,13 +127,13 @@ def test_equivariant_power_spectrum_neighbors_to_properties():
     )
 
     # Permute properties dimensions to match ``powspec_1`` and sort
-    powspec_2 = metatensor.sort(
-        metatensor.permute_dimensions(powspec_2, "properties", [2, 4, 0, 1, 3, 5])
+    powspec_2 = mts.sort(
+        mts.permute_dimensions(powspec_2, "properties", [2, 4, 0, 1, 3, 5])
     )
 
-    # Check equivalent
-    metatensor.equal_metadata_raise(powspec_1, powspec_2)
-    metatensor.equal_raise(powspec_1, powspec_2)
+    # Check equivalence
+    mts.equal_metadata_raise(powspec_1, powspec_2)
+    mts.equal_raise(powspec_1, powspec_2)
 
 
 def test_sample_selection() -> None:
@@ -147,11 +147,9 @@ def test_sample_selection() -> None:
 
     powspec_calc = EquivariantPowerSpectrum(SphericalExpansion(**SPHEX_HYPERS_SMALL))
 
-    label_1st = metatensor.Labels(
-        ["system", "atom"], np.array([[0, 0]], dtype=np.int32)
-    )
+    label_1st = mts.Labels(["system", "atom"], np.array([[0, 0]], dtype=np.int32))
 
-    label_2nd = metatensor.Labels(
+    label_2nd = mts.Labels(
         ["system", "atom"], np.array([[0, 1], [0, 2]], dtype=np.int32)
     )
 
@@ -163,12 +161,12 @@ def test_sample_selection() -> None:
         frame, neighbors_to_properties=True, selected_samples=label_2nd
     )
 
-    powspec_3 = metatensor.join([powspec_1, powspec_2], axis="samples")
+    powspec_3 = mts.join([powspec_1, powspec_2], axis="samples")
     powspec_4 = powspec_calc.compute(frame, neighbors_to_properties=True)
 
-    assert metatensor.equal(powspec_3, powspec_4)
-    assert not metatensor.equal(powspec_2, powspec_4)
-    assert not metatensor.equal(powspec_1, powspec_4)
+    assert mts.equal(powspec_3, powspec_4)
+    assert not mts.equal(powspec_2, powspec_4)
+    assert not mts.equal(powspec_1, powspec_4)
 
 
 def test_fill_types_option() -> None:
diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py
index cbb66a188..2c5ba4c36 100644
--- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py
+++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py
@@ -1,4 +1,4 @@
-import metatensor
+import metatensor as mts
 import numpy as np
 import pytest
 from metatensor import Labels
@@ -92,23 +92,23 @@ def test_equivariant_power_spectrum_vs_equivariant_power_spectrum_by_pair():
     )
 
     # Manipulate metadata to match
-    reduced_ps_by_pair = metatensor.rename_dimension(
+    reduced_ps_by_pair = mts.rename_dimension(
         ps_by_pair, "keys", "first_atom_type", "center_type"
     )
-    reduced_ps_by_pair = metatensor.rename_dimension(
+    reduced_ps_by_pair = mts.rename_dimension(
         reduced_ps_by_pair,
         "keys",
         "second_atom_type",
         "neighbor_2_type",
     )
-    reduced_ps_by_pair = metatensor.rename_dimension(
+    reduced_ps_by_pair = mts.rename_dimension(
         reduced_ps_by_pair,
         "samples",
         "first_atom",
         "atom",
     )
     # Sum over `second_atom` and `cell_shift` samples
-    reduced_ps_by_pair = metatensor.sum_over_samples(
+    reduced_ps_by_pair = mts.sum_over_samples(
         reduced_ps_by_pair,
         ["second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"],
     )
@@ -136,11 +136,9 @@ def test_sample_selection() -> None:
         SphericalExpansionByPair(**SPHEX_HYPERS_SMALL),
     )
 
-    label_1st = metatensor.Labels(
-        ["system", "atom"], np.array([[0, 0]], dtype=np.int32)
-    )
+    label_1st = mts.Labels(["system", "atom"], np.array([[0, 0]], dtype=np.int32))
 
-    label_2nd = metatensor.Labels(
+    label_2nd = mts.Labels(
         ["system", "atom"], np.array([[0, 1], [0, 2]], dtype=np.int32)
     )
 
@@ -152,12 +150,12 @@ def test_sample_selection() -> None:
         frame, neighbors_to_properties=True, selected_samples=label_2nd
     )
 
-    powspec_3 = metatensor.join([powspec_1, powspec_2], axis="samples")
+    powspec_3 = mts.join([powspec_1, powspec_2], axis="samples")
     powspec_4 = powspec_by_pair_calc.compute(frame, neighbors_to_properties=True)
 
-    assert metatensor.equal(powspec_3, powspec_4)
-    assert not metatensor.equal(powspec_2, powspec_4)
-    assert not metatensor.equal(powspec_1, powspec_4)
+    assert mts.equal(powspec_3, powspec_4)
+    assert not mts.equal(powspec_2, powspec_4)
+    assert not mts.equal(powspec_1, powspec_4)
 
 
 def test_equivariant_power_spectrum_neighbors_to_properties():
@@ -185,14 +183,14 @@ def test_equivariant_power_spectrum_neighbors_to_properties():
     )
 
     # Permute properties dimensions to match ``powspec_1`` and sort
-    powspec_2 = metatensor.sort(
-        metatensor.permute_dimensions(powspec_2, "properties", [2, 0, 1, 3, 4])
+    powspec_2 = mts.sort(
+        mts.permute_dimensions(powspec_2, "properties", [2, 0, 1, 3, 4])
     )
 
     # Check equivalent
-    powspec_1 = metatensor.sort(powspec_1)
-    metatensor.equal_metadata_raise(powspec_1, powspec_2)
-    metatensor.allclose_raise(powspec_1, powspec_2)
+    powspec_1 = mts.sort(powspec_1)
+    mts.equal_metadata_raise(powspec_1, powspec_2)
+    mts.allclose_raise(powspec_1, powspec_2)
 
 
 def test_fill_types_option() -> None:
diff --git a/python/featomic_torch/tests/clebsch_gordan/density_correlations.py b/python/featomic_torch/tests/clebsch_gordan/density_correlations.py
index 37059f731..93fab33cd 100644
--- a/python/featomic_torch/tests/clebsch_gordan/density_correlations.py
+++ b/python/featomic_torch/tests/clebsch_gordan/density_correlations.py
@@ -2,7 +2,7 @@
 import io
 import os
 
-import metatensor.torch
+import metatensor.torch as mts
 import pytest
 import torch
 from metatensor.torch import Labels
@@ -71,8 +71,8 @@ def test_torch_script_correlate_density_angular_selection(
     # Compute the reference and scripted results
     ref_nu_2 = calculator.compute(nu_1, selected_keys=selected_keys)
     scripted_nu_2 = scripted_calculator.compute(nu_1, selected_keys=selected_keys)
-    metatensor.torch.equal_metadata_raise(scripted_nu_2, ref_nu_2)
-    assert metatensor.torch.allclose(scripted_nu_2, ref_nu_2)
+    mts.equal_metadata_raise(scripted_nu_2, ref_nu_2)
+    assert mts.allclose(scripted_nu_2, ref_nu_2)
 
 
 @pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])

From 34bebf2dfce4172b8cde4170d02ae6469af47829 Mon Sep 17 00:00:00 2001
From: Guillaume Fraux <guillaume.fraux@epfl.ch>
Date: Fri, 14 Nov 2025 17:23:46 +0100
Subject: [PATCH 3/3] General code cleanup

---
 python/featomic/examples/keys-selection.py    |  6 +-
 .../tests/calculators/keys_selection.py       | 60 ++++---------------
 .../tests/calculators/properties_selection.py | 27 +++------
 .../tests/calculators/sample_selection.py     | 33 +++-------
 .../equivariant_power_spectrum.py             |  6 +-
 .../equivariant_power_spectrum_by_pair.py     |  6 +-
 .../featomic/torch/clebsch_gordan.py          |  2 +-
 python/featomic_torch/tests/export.py         |  5 +-
 8 files changed, 35 insertions(+), 110 deletions(-)

diff --git a/python/featomic/examples/keys-selection.py b/python/featomic/examples/keys-selection.py
index f323d5f97..37fc1a82f 100644
--- a/python/featomic/examples/keys-selection.py
+++ b/python/featomic/examples/keys-selection.py
@@ -57,7 +57,7 @@
 
 selection = Labels(
     names=["center_type", "neighbor_1_type", "neighbor_2_type"],
-    values=np.array([[1, 1, 1], [4, 4, 4]], dtype=np.int32),
+    values=np.array([[1, 1, 1], [4, 4, 4]]),
 )
 selected_descriptor = calculator.compute(frames, selected_keys=selection)
 
@@ -104,7 +104,7 @@
 
 keys = Labels(
     names=["center_type", "neighbor_1_type", "neighbor_2_type"],
-    values=np.array([[1, 1, 1], [8, 8, 8]], dtype=np.int32),
+    values=np.array([[1, 1, 1], [8, 8, 8]]),
 )
 
 selected_properties = TensorMap(keys, blocks)
@@ -116,7 +116,7 @@
 
 selected_keys = Labels(
     names=["center_type", "neighbor_1_type", "neighbor_2_type"],
-    values=np.array([[1, 1, 1]], dtype=np.int32),
+    values=np.array([[1, 1, 1]]),
 )
 
 descriptor = calculator.compute(
diff --git a/python/featomic/tests/calculators/keys_selection.py b/python/featomic/tests/calculators/keys_selection.py
index 7499926a1..c236cf3f5 100644
--- a/python/featomic/tests/calculators/keys_selection.py
+++ b/python/featomic/tests/calculators/keys_selection.py
@@ -44,10 +44,7 @@ def test_selection_existing():
     assert descriptor.keys.values.tolist() == [[1], [8]]
 
     # Manually select the keys
-    selected_keys = Labels(
-        names=["center_type"],
-        values=np.array([[1]], dtype=np.int32),
-    )
+    selected_keys = Labels(names=["center_type"], values=np.array([[1]]))
     descriptor = calculator.compute(system, selected_keys=selected_keys)
 
     assert len(descriptor.keys) == 1
@@ -73,10 +70,7 @@ def test_selection_partial():
     )
 
     # Manually select the keys
-    selected_keys = Labels(
-        names=["center_type"],
-        values=np.array([[1]], dtype=np.int32),
-    )
+    selected_keys = Labels(names=["center_type"], values=np.array([[1]]))
     descriptor = calculator.compute(system, selected_keys=selected_keys)
 
     assert len(descriptor.keys) == 4
@@ -99,10 +93,7 @@ def test_select_key_not_in_systems():
     calculator = DummyCalculator(cutoff=3.2, delta=2, name="")
 
     # Manually select the keys
-    selected_keys = Labels(
-        names=["center_type"],
-        values=np.array([[4]], dtype=np.int32),
-    )
+    selected_keys = Labels(names=["center_type"], values=np.array([[4]]))
     descriptor = calculator.compute(system, selected_keys=selected_keys)
 
     C_block = descriptor.block(center_type=4)
@@ -113,26 +104,14 @@ def test_predefined_selection():
     system = SystemForTests()
     calculator = DummyCalculator(cutoff=3.2, delta=2, name="")
 
-    selected_keys = Labels(
-        names=["center_type"],
-        values=np.array([[1]], dtype=np.int32),
-    )
+    selected_keys = Labels(names=["center_type"], values=np.array([[1]]))
 
-    keys = Labels(
-        names=["center_type"],
-        values=np.array([[1], [8]], dtype=np.int32),
-    )
+    keys = Labels(names=["center_type"], values=np.array([[1], [8]]))
 
     # selection from TensorMap
     selected = [
-        Labels(
-            names=["index_delta", "x_y_z"],
-            values=np.array([[1, 0]], dtype=np.int32),
-        ),
-        Labels(
-            names=["index_delta", "x_y_z"],
-            values=np.array([[0, 1]], dtype=np.int32),
-        ),
+        Labels(names=["index_delta", "x_y_z"], values=np.array([[1, 0]])),
+        Labels(names=["index_delta", "x_y_z"], values=np.array([[0, 1]])),
     ]
     selected_properties = _tensor_map_selection("properties", keys, selected)
 
@@ -153,10 +132,7 @@ def test_name_errors():
     system = SystemForTests()
     calculator = DummyCalculator(cutoff=3.2, delta=2, name="")
 
-    selected_keys = Labels(
-        names=["bad_name"],
-        values=np.array([0, 3, 1], dtype=np.int32).reshape(3, 1),
-    )
+    selected_keys = Labels(names=["bad_name"], values=np.array([[0], [3], [1]]))
 
     message = (
         "invalid parameter: 'bad_name' in keys selection is not "
@@ -182,25 +158,13 @@ def test_key_errors():
     # in the case where both selected_properties/selected_samples and
     # selected_keys are given, the selected keys must be in the keys of the
     # predefined tensor_map
-    selected_keys = Labels(
-        names=["center_type"],
-        values=np.array([[4]], dtype=np.int32),
-    )
+    selected_keys = Labels(names=["center_type"], values=np.array([[4]]))
 
-    keys = Labels(
-        names=["center_type"],
-        values=np.array([[1], [8]], dtype=np.int32),
-    )
+    keys = Labels(names=["center_type"], values=np.array([[1], [8]]))
 
     selected = [
-        Labels(
-            names=["index_delta", "x_y_z"],
-            values=np.array([[1, 0]], dtype=np.int32),
-        ),
-        Labels(
-            names=["index_delta", "x_y_z"],
-            values=np.array([[0, 1]], dtype=np.int32),
-        ),
+        Labels(names=["index_delta", "x_y_z"], values=np.array([[1, 0]])),
+        Labels(names=["index_delta", "x_y_z"], values=np.array([[0, 1]])),
     ]
     selected_properties = _tensor_map_selection("properties", keys, selected)
 
diff --git a/python/featomic/tests/calculators/properties_selection.py b/python/featomic/tests/calculators/properties_selection.py
index 61bf8a216..96d31060d 100644
--- a/python/featomic/tests/calculators/properties_selection.py
+++ b/python/featomic/tests/calculators/properties_selection.py
@@ -41,7 +41,7 @@ def test_selection():
     # Manually constructing the selected properties
     selected_properties = Labels(
         names=["index_delta", "x_y_z"],
-        values=np.array([[1, 0]], dtype=np.int32),
+        values=np.array([[1, 0]]),
     )
     descriptor = calculator.compute(
         system, use_native_system=False, selected_properties=selected_properties
@@ -65,7 +65,7 @@ def test_subset_variables():
     # Only a subset of the variables defined
     selected_properties = Labels(
         names=["index_delta"],
-        values=np.array([[1]], dtype=np.int32),
+        values=np.array([[1]]),
     )
     descriptor = calculator.compute(
         system, use_native_system=False, selected_properties=selected_properties
@@ -88,8 +88,7 @@ def test_empty_selection():
 
     # empty selected features
     selected_properties = Labels(
-        names=["index_delta", "x_y_z"],
-        values=np.array([], dtype=np.int32).reshape(0, 2),
+        names=["index_delta", "x_y_z"], values=np.empty((0, 2), dtype=np.int32)
     )
     descriptor = calculator.compute(
         system, use_native_system=False, selected_properties=selected_properties
@@ -106,21 +105,12 @@ def test_predefined_selection():
     system = SystemForTests()
     calculator = DummyCalculator(cutoff=3.2, delta=2, name="")
 
-    keys = Labels(
-        names=["center_type"],
-        values=np.array([[1], [8]], dtype=np.int32),
-    )
+    keys = Labels(names=["center_type"], values=np.array([[1], [8]]))
 
     # selection from TensorMap
     selected = [
-        Labels(
-            names=["index_delta", "x_y_z"],
-            values=np.array([[1, 0]], dtype=np.int32),
-        ),
-        Labels(
-            names=["index_delta", "x_y_z"],
-            values=np.array([[0, 1]], dtype=np.int32),
-        ),
+        Labels(names=["index_delta", "x_y_z"], values=np.array([[1, 0]])),
+        Labels(names=["index_delta", "x_y_z"], values=np.array([[0, 1]])),
     ]
     selected_properties = _tensor_map_selection("properties", keys, selected)
 
@@ -143,10 +133,7 @@ def test_errors():
     system = SystemForTests()
     calculator = DummyCalculator(cutoff=3.2, delta=2, name="")
 
-    selected_properties = Labels(
-        names=["bad_name"],
-        values=np.array([0, 3, 1], dtype=np.int32).reshape(3, 1),
-    )
+    selected_properties = Labels(names=["bad_name"], values=np.array([[0], [3], [1]]))
 
     message = (
         "invalid parameter: 'bad_name' in properties selection is not "
diff --git a/python/featomic/tests/calculators/sample_selection.py b/python/featomic/tests/calculators/sample_selection.py
index 3bee6fce1..131fa1498 100644
--- a/python/featomic/tests/calculators/sample_selection.py
+++ b/python/featomic/tests/calculators/sample_selection.py
@@ -40,8 +40,7 @@ def test_selection():
 
     # Manually constructing the selected samples
     selected_samples = Labels(
-        names=["system", "atom"],
-        values=np.array([(0, 0), (0, 3), (0, 1)], dtype=np.int32),
+        names=["system", "atom"], values=np.array([[0, 0], [0, 3], [0, 1]])
     )
     descriptor = calculator.compute(
         system, use_native_system=False, selected_samples=selected_samples
@@ -62,10 +61,7 @@ def test_subset_variables():
     calculator = DummyCalculator(cutoff=3.2, delta=2, name="")
 
     # Only a subset of the variables defined
-    selected_samples = Labels(
-        names=["atom"],
-        values=np.array([0, 3, 1], dtype=np.int32).reshape(3, 1),
-    )
+    selected_samples = Labels(names=["atom"], values=np.array([0, 3, 1]).reshape(3, 1))
     descriptor = calculator.compute(
         system, use_native_system=False, selected_samples=selected_samples
     )
@@ -85,10 +81,7 @@ def test_empty_selection():
     calculator = DummyCalculator(cutoff=3.2, delta=2, name="")
 
     # empty selected samples
-    selected_samples = Labels(
-        names=["atom"],
-        values=np.empty((0, 1), dtype=np.int32),
-    )
+    selected_samples = Labels(names=["atom"], values=np.empty((0, 1), dtype=np.int32))
     descriptor = calculator.compute(
         system, use_native_system=False, selected_samples=selected_samples
     )
@@ -104,21 +97,12 @@ def test_predefined_selection():
     system = SystemForTests()
     calculator = DummyCalculator(cutoff=3.2, delta=2, name="")
 
-    keys = Labels(
-        names=["center_type"],
-        values=np.array([[1], [8]], dtype=np.int32),
-    )
+    keys = Labels(names=["center_type"], values=np.array([[1], [8]]))
 
     # selection from TensorMap
     selected = [
-        Labels(
-            names=["system", "atom"],
-            values=np.array([[0, 1]], dtype=np.int32),
-        ),
-        Labels(
-            names=["system", "atom"],
-            values=np.array([[0, 3]], dtype=np.int32),
-        ),
+        Labels(names=["system", "atom"], values=np.array([[0, 1]])),
+        Labels(names=["system", "atom"], values=np.array([[0, 3]])),
     ]
     selected_samples = _tensor_map_selection("samples", keys, selected)
 
@@ -139,10 +123,7 @@ def test_errors():
     system = SystemForTests()
     calculator = DummyCalculator(cutoff=3.2, delta=2, name="")
 
-    samples = Labels(
-        names=["bad_name"],
-        values=np.array([0, 3, 1], dtype=np.int32).reshape(3, 1),
-    )
+    samples = Labels(names=["bad_name"], values=np.array([0, 3, 1]).reshape(3, 1))
 
     message = (
         "invalid parameter: 'bad_name' in samples selection is not part "
diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py
index d73fa55b0..e4c37f153 100644
--- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py
+++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py
@@ -147,11 +147,9 @@ def test_sample_selection() -> None:
 
     powspec_calc = EquivariantPowerSpectrum(SphericalExpansion(**SPHEX_HYPERS_SMALL))
 
-    label_1st = mts.Labels(["system", "atom"], np.array([[0, 0]], dtype=np.int32))
+    label_1st = mts.Labels(["system", "atom"], np.array([[0, 0]]))
 
-    label_2nd = mts.Labels(
-        ["system", "atom"], np.array([[0, 1], [0, 2]], dtype=np.int32)
-    )
+    label_2nd = mts.Labels(["system", "atom"], np.array([[0, 1], [0, 2]]))
 
     powspec_1 = powspec_calc.compute(
         frame, neighbors_to_properties=True, selected_samples=label_1st
diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py
index 2c5ba4c36..716a0eddf 100644
--- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py
+++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py
@@ -136,11 +136,9 @@ def test_sample_selection() -> None:
         SphericalExpansionByPair(**SPHEX_HYPERS_SMALL),
     )
 
-    label_1st = mts.Labels(["system", "atom"], np.array([[0, 0]], dtype=np.int32))
+    label_1st = mts.Labels(["system", "atom"], np.array([[0, 0]]))
 
-    label_2nd = mts.Labels(
-        ["system", "atom"], np.array([[0, 1], [0, 2]], dtype=np.int32)
-    )
+    label_2nd = mts.Labels(["system", "atom"], np.array([[0, 1], [0, 2]]))
 
     powspec_1 = powspec_by_pair_calc.compute(
         frame, neighbors_to_properties=True, selected_samples=label_1st
diff --git a/python/featomic_torch/featomic/torch/clebsch_gordan.py b/python/featomic_torch/featomic/torch/clebsch_gordan.py
index a9aa2bc78..d282ec8a3 100644
--- a/python/featomic_torch/featomic/torch/clebsch_gordan.py
+++ b/python/featomic_torch/featomic/torch/clebsch_gordan.py
@@ -6,7 +6,7 @@
 import metatensor.torch
 import torch
 
-import featomic.utils
+import featomic.clebsch_gordan
 
 from .calculator_base import CalculatorModule
 from .system import System
diff --git a/python/featomic_torch/tests/export.py b/python/featomic_torch/tests/export.py
index f38db35f7..2eb9a6eea 100644
--- a/python/featomic_torch/tests/export.py
+++ b/python/featomic_torch/tests/export.py
@@ -75,10 +75,7 @@ def forward(
             samples = soap.block().samples
         else:
             features = soap.block().values.sum(dim=0, keepdim=True)
-            samples = Labels(
-                ["system"],
-                torch.arange(len(systems), dtype=torch.int32).reshape(-1, 1),
-            )
+            samples = Labels(["system"], torch.arange(len(systems)).reshape(-1, 1))
 
         block = TensorBlock(
             values=self.linear(features),