diff --git a/docs/src/references/api/torch/index.rst b/docs/src/references/api/torch/index.rst
index 43b730c66..401dc70b7 100644
--- a/docs/src/references/api/torch/index.rst
+++ b/docs/src/references/api/torch/index.rst
@@ -27,6 +27,7 @@ of rascaline are documented below for an usage from Python:
 
     systems
     calculators
+    utils/index
 
 --------------------------------------------------------------------------------
 
diff --git a/docs/src/references/api/torch/utils/index.rst b/docs/src/references/api/torch/utils/index.rst
new file mode 100644
index 000000000..1351283df
--- /dev/null
+++ b/docs/src/references/api/torch/utils/index.rst
@@ -0,0 +1,10 @@
+Utils
+=====
+
+Utility functions and classes that extend the core usage of rascaline-torch
+
+
+.. toctree::
+    :maxdepth: 1
+
+    power-spectrum
diff --git a/docs/src/references/api/torch/utils/power-spectrum.rst b/docs/src/references/api/torch/utils/power-spectrum.rst
new file mode 100644
index 000000000..12633c552
--- /dev/null
+++ b/docs/src/references/api/torch/utils/power-spectrum.rst
@@ -0,0 +1,6 @@
+PowerSpectrum
+=============
+
+.. autoclass:: rascaline.torch.utils.PowerSpectrum
+    :members:
+    :show-inheritance:
diff --git a/python/rascaline-torch/rascaline/torch/utils/__init__.py b/python/rascaline-torch/rascaline/torch/utils/__init__.py
new file mode 100644
index 000000000..bc31628b9
--- /dev/null
+++ b/python/rascaline-torch/rascaline/torch/utils/__init__.py
@@ -0,0 +1,4 @@
+from .power_spectrum import PowerSpectrum
+
+
+__all__ = ["PowerSpectrum"]
diff --git a/python/rascaline-torch/rascaline/torch/utils/_classes.py b/python/rascaline-torch/rascaline/torch/utils/_classes.py
new file mode 100644
index 000000000..55a43e87e
--- /dev/null
+++ b/python/rascaline-torch/rascaline/torch/utils/_classes.py
@@ -0,0 +1,18 @@
+from equistore.torch import Labels, TensorBlock, TensorMap
+from torch.nn import Module as torch_nn_module
+
+from rascaline.utils import _dispatch
+
+from ..calculator_base import CalculatorModule as CalculatorBase
+from ..system import System as IntoSystem
+
+
+__all__ = [
+    "_dispatch",
+    "CalculatorBase",
+    "IntoSystem",
+    "Labels",
+    "torch_nn_module",
+    "TensorBlock",
+    "TensorMap",
+]
diff --git a/python/rascaline-torch/rascaline/torch/utils/power_spectrum.py b/python/rascaline-torch/rascaline/torch/utils/power_spectrum.py
new file mode 100644
index 000000000..b13660d49
--- /dev/null
+++ b/python/rascaline-torch/rascaline/torch/utils/power_spectrum.py
@@ -0,0 +1,20 @@
+import importlib
+import sys
+
+import rascaline.utils.power_spectrum
+
+
+# For details what is happening here take a look an `rascaline.torch.calculators`.
+spec = importlib.util.spec_from_file_location(
+    # create a module with this name
+    "rascaline.torch.utils.power_spectrum",
+    # using the code from there
+    rascaline.utils.power_spectrum.__file__,
+)
+module = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = module
+spec.loader.exec_module(module)
+
+# don't forget to also update `rascaline/utils/__init__.py` and
+# `rascaline/torch/utils.__init__.py` when modifying this file
+PowerSpectrum = module.PowerSpectrum
diff --git a/python/rascaline-torch/tests/utils/power_spectrum.py b/python/rascaline-torch/tests/utils/power_spectrum.py
new file mode 100644
index 000000000..631d0db01
--- /dev/null
+++ b/python/rascaline-torch/tests/utils/power_spectrum.py
@@ -0,0 +1,51 @@
+import torch
+from packaging import version
+
+from rascaline.torch import System
+from rascaline.torch.calculators import SphericalExpansion
+from rascaline.torch.utils import PowerSpectrum
+
+
+def system():
+    return System(
+        species=torch.tensor([1, 1, 8, 8]),
+        positions=torch.tensor([[0.0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]),
+        cell=torch.tensor([[10, 0, 0], [0, 10, 0], [0, 0, 10]]),
+    )
+
+
+def calculator():
+    return SphericalExpansion(
+        cutoff=5.0,
+        max_radial=6,
+        max_angular=4,
+        atomic_gaussian_width=0.3,
+        center_atom_weight=1.0,
+        radial_basis={
+            "Gto": {},
+        },
+        cutoff_function={
+            "ShiftedCosine": {"width": 0.5},
+        },
+    )
+
+
+def check_operation(PowerSpectrum):
+    # this only runs basic checks functionality checks, and that the code produces
+    # output with the right type
+
+    descriptor = calculator().compute(system(), gradients=["positions"])
+
+    assert isinstance(descriptor, torch.ScriptObject)
+    if version.parse(torch.__version__) >= version.parse("2.1"):
+        assert descriptor._type().name() == "TensorMap"
+
+
+def test_operation_as_python():
+    check_operation(PowerSpectrum)
+
+
+def test_operation_as_torch_script():
+    scripted = torch.jit.script(calculator())
+
+    check_operation(scripted)
diff --git a/python/rascaline/rascaline/utils/_classes.py b/python/rascaline/rascaline/utils/_classes.py
new file mode 100644
index 000000000..d20302569
--- /dev/null
+++ b/python/rascaline/rascaline/utils/_classes.py
@@ -0,0 +1,20 @@
+from equistore.core import Labels, TensorBlock, TensorMap
+
+from ..calculator_base import CalculatorBase
+from ..systems import IntoSystem
+from . import _dispatch
+
+
+# dummy object which is only relevant for torch
+torch_nn_module = object
+
+
+__all__ = [
+    "CalculatorBase",
+    "IntoSystem",
+    "Labels",
+    "TensorBlock",
+    "TensorMap",
+    "_dispatch",
+    "torch_nn_module",
+]
diff --git a/python/rascaline/rascaline/utils/_dispatch.py b/python/rascaline/rascaline/utils/_dispatch.py
new file mode 100644
index 000000000..48cfd2bee
--- /dev/null
+++ b/python/rascaline/rascaline/utils/_dispatch.py
@@ -0,0 +1,103 @@
+"""Helper functions to dispatch methods between numpy and torch.
+
+The functions are similar to those in equistore-operations. Missing functions may
+already exist there. Functions are ordered alphabetically.
+"""
+
+from typing import List
+
+import numpy as np
+
+
+try:
+    import torch
+    from torch import Tensor as TorchTensor
+except ImportError:
+
+    class TorchTensor:
+        pass
+
+
+UNKNOWN_ARRAY_TYPE = (
+    "unknown array type, only numpy arrays and torch tensors are supported"
+)
+
+
+def _check_all_np_ndarray(arrays):
+    for array in arrays:
+        if not isinstance(array, np.ndarray):
+            raise TypeError(
+                f"expected argument to be a np.ndarray, but got {type(array)}"
+            )
+
+
+def _check_all_torch_tensor(arrays):
+    for array in arrays:
+        if not isinstance(array, TorchTensor):
+            raise TypeError(
+                f"expected argument to be a torch.Tensor, but got {type(array)}"
+            )
+
+
+def list_to_array(array, data: List):
+    """Create an object from data with the same type as ``array``."""
+    if isinstance(array, TorchTensor):
+        return torch.Tensor(data, device=array.device, dtype=array.dtype)
+    elif isinstance(array, np.ndarray):
+        return np.array(data, dtype=array.dtype)
+    else:
+        raise TypeError(UNKNOWN_ARRAY_TYPE)
+
+
+def matmul(a, b, out=None):
+    """Matrix product of two arrays."""
+    if isinstance(a, TorchTensor):
+        _check_all_torch_tensor([b])
+        return torch.matmul(a, b, out=out)
+    elif isinstance(a, np.ndarray):
+        _check_all_np_ndarray([b])
+        return np.matmul(a, b, out=out)
+    else:
+        raise TypeError(UNKNOWN_ARRAY_TYPE)
+
+
+def unique(array, return_inverse=False, return_counts=False, axis=None):
+    """Find the unique elements of an array."""
+    if isinstance(array, TorchTensor):
+        return torch.unique(
+            array, return_inverse=return_inverse, return_counts=return_counts, dim=axis
+        )
+    elif isinstance(array, np.ndarray):
+        return np.unique(
+            array, return_inverse=return_inverse, return_counts=return_counts, axis=axis
+        )
+    else:
+        raise TypeError(UNKNOWN_ARRAY_TYPE)
+
+
+def zeros_like(array, shape=None, requires_grad=False):
+    """
+    Create an array filled with zeros, with the given ``shape``, and similar
+    dtype, device and other options as ``array``.
+
+    If ``shape`` is :py:obj:`None`, the array shape is used instead.
+    ``requires_grad`` is only used for torch tensors, and set the corresponding
+    value on the returned array.
+
+    This is the equivalent to ``np.zeros_like(array, shape=shape)``.
+    """
+    if isinstance(array, TorchTensor):
+        if shape is None:
+            shape = array.size()
+
+        return torch.zeros(
+            shape,
+            dtype=array.dtype,
+            layout=array.layout,
+            device=array.device,
+            requires_grad=requires_grad,
+        )
+    elif isinstance(array, np.ndarray):
+        return np.zeros_like(array, shape=shape, subok=False)
+    else:
+        raise TypeError(UNKNOWN_ARRAY_TYPE)
diff --git a/python/rascaline/rascaline/utils/power_spectrum.py b/python/rascaline/rascaline/utils/power_spectrum.py
index 79949cca8..dc594d1f6 100644
--- a/python/rascaline/rascaline/utils/power_spectrum.py
+++ b/python/rascaline/rascaline/utils/power_spectrum.py
@@ -2,14 +2,18 @@
 from math import sqrt
 from typing import List, Optional, Union
 
-import numpy as np
-from equistore.core import Labels, TensorBlock, TensorMap
-
-from ..calculators import CalculatorBase
-from ..systems import IntoSystem
-
-
-class PowerSpectrum:
+from ._classes import (
+    CalculatorBase,
+    IntoSystem,
+    Labels,
+    TensorBlock,
+    TensorMap,
+    _dispatch,
+    torch_nn_module,
+)
+
+
+class PowerSpectrum(torch_nn_module):
     def __init__(
         self,
         calculator_1: CalculatorBase,
@@ -119,6 +123,8 @@ def __init__(
             If you are interested in the SOAP power spectrum you can the use the
             faster :py:class:`rascaline.SoapPowerSpectrum`.
         """  # noqa E501
+        super().__init__()
+
         self.calculator_1 = calculator_1
         self.calculator_2 = calculator_2
 
@@ -150,7 +156,6 @@ def name(self):
     def compute(
         self,
         systems: Union[IntoSystem, List[IntoSystem]],
-        *,
         gradients: Optional[List[str]] = None,
         use_native_system: bool = True,
     ) -> TensorMap:
@@ -171,9 +176,7 @@ def compute(
                     )
 
         spherical_expansion_1 = self.calculator_1.compute(
-            systems=systems,
-            gradients=gradients,
-            use_native_system=use_native_system,
+            systems=systems, gradients=gradients, use_native_system=use_native_system
         )
 
         expected_key_names = [
@@ -188,9 +191,9 @@ def compute(
         # merging blocks along the ``sample`` direction might be not possible.
         keys_to_move = Labels(
             names="species_neighbor",
-            values=np.unique(spherical_expansion_1.keys["species_neighbor"]).reshape(
-                -1, 1
-            ),
+            values=_dispatch.unique(
+                spherical_expansion_1.keys["species_neighbor"]
+            ).reshape(-1, 1),
         )
 
         spherical_expansion_1 = spherical_expansion_1.keys_to_properties(keys_to_move)
@@ -208,7 +211,7 @@ def compute(
 
             keys_to_move = Labels(
                 names="species_neighbor",
-                values=np.unique(
+                values=_dispatch.unique(
                     spherical_expansion_2.keys["species_neighbor"]
                 ).reshape(-1, 1),
             )
@@ -224,28 +227,35 @@ def compute(
             factor = 1 / sqrt(2 * ell + 1)
             # Find that block indices that have the same spherical_harmonics_l and
             # species_center
-            blocks_2 = spherical_expansion_2.blocks(
-                spherical_harmonics_l=ell, species_center=species_center
+            selection = Labels(
+                names=["spherical_harmonics_l", "species_center"],
+                values=_dispatch.list_to_array(
+                    array=spherical_expansion_1.keys.values,
+                    data=[[ell, species_center]],
+                ),
             )
+            blocks_2 = spherical_expansion_2.blocks(selection)
             for block_2 in blocks_2:
                 # Make sure that samples are the same. This should not happen.
                 assert block_1.samples == block_2.samples
 
+                data = []
+                for properties_1 in block_1.properties.values:
+                    for properties_2 in block_2.properties.values:
+                        data.append(properties_1.tolist() + properties_2.tolist())
+
                 properties = Labels(
                     names=["species_neighbor_1", "n1", "species_neighbor_2", "n2"],
-                    values=np.array(
-                        [
-                            properties_1.tolist() + properties_2.tolist()
-                            for properties_1 in block_1.properties.values
-                            for properties_2 in block_2.properties.values
-                        ],
-                        dtype=np.int32,
+                    values=_dispatch.list_to_array(
+                        array=block_1.properties.values, data=data
                     ),
                 )
 
                 # Compute the invariants by summation and store the results this is
                 # equivalent to an einsum with: ima, imb -> iab
-                data = factor * np.matmul(block_1.values.swapaxes(1, 2), block_2.values)
+                data = factor * _dispatch.matmul(
+                    block_1.values.swapaxes(1, 2), block_2.values
+                )
 
                 new_block = TensorBlock(
                     values=data.reshape(data.shape[0], -1),
@@ -263,11 +273,25 @@ def compute(
 
         keys = Labels(
             names=["l", "species_center"],
-            values=np.array(keys, dtype=np.int32),
+            values=_dispatch.list_to_array(
+                array=spherical_expansion_1.keys.values, data=keys
+            ),
         )
 
         return TensorMap(keys, blocks).keys_to_properties("l")
 
+    def forward(
+        self,
+        systems: Union[IntoSystem, List[IntoSystem]],
+        gradients: Optional[List[str]] = None,
+        use_native_system: bool = True,
+    ) -> TensorMap:
+        """forward just calls :py:meth:`PowerSpectrum.compute`"""
+
+        return self.compute(
+            systems=systems, gradients=gradients, use_native_system=use_native_system
+        )
+
 
 def _positions_gradients(new_block, block_1, block_2, factor):
     gradient_1 = block_1.gradient("positions")
@@ -275,7 +299,9 @@ def _positions_gradients(new_block, block_1, block_2, factor):
 
     if len(gradient_1.samples) == 0 or len(gradient_2.samples) == 0:
         gradients_samples = Labels.empty(names=["sample", "structure", "atom"])
-        gradient_values = np.array([]).reshape(0, 1, len(new_block.properties))
+        gradient_values = _dispatch.list_to_array(
+            array=gradient_1.values, data=[]
+        ).reshape(0, 1, len(new_block.properties))
     else:
         # The "sample" dimension in the power spectrum gradient samples do
         # not necessarily matches the "sample" dimension in the spherical
@@ -287,14 +313,17 @@ def _positions_gradients(new_block, block_1, block_2, factor):
             grad2_sample_idxs,
         ) = gradient_1.samples.union_and_mapping(gradient_2.samples)
 
-        gradient_values = np.zeros(
-            [gradients_samples.values.shape[0], 3, len(new_block.properties)]
+        gradient_values = _dispatch.zeros_like(
+            array=gradient_1.values,
+            shape=[gradients_samples.values.shape[0], 3, len(new_block.properties)],
         )
 
         # this is equivalent to an einsum with: ixma, imb -> ixab
-        gradient_1_values = factor * np.matmul(
+        block_2_values = block_2.values[gradient_1.samples["sample"]]
+        new_shape = block_2_values.shape[:1] + (-1,) + block_2_values.shape[1:]
+        gradient_1_values = factor * _dispatch.matmul(
             gradient_1.values.swapaxes(2, 3),
-            block_2.values[gradient_1.samples["sample"], np.newaxis],
+            block_2_values.reshape(new_shape),
         )
 
         gradient_values[grad1_sample_idxs] += gradient_1_values.reshape(
@@ -302,8 +331,10 @@ def _positions_gradients(new_block, block_1, block_2, factor):
         )
 
         # this is equivalent to an einsum with: ima, ixmb -> ixab
-        gradient_values_2 = factor * np.matmul(
-            block_1.values[gradient_2.samples["sample"], np.newaxis].swapaxes(2, 3),
+        block_1_values = block_1.values[gradient_2.samples["sample"]]
+        new_shape = block_1_values.shape[:1] + (-1,) + block_1_values.shape[1:]
+        gradient_values_2 = factor * _dispatch.matmul(
+            block_1_values.reshape(new_shape).swapaxes(2, 3),
             gradient_2.values,
         )
 
diff --git a/python/rascaline/tests/utils/power_spectrum.py b/python/rascaline/tests/utils/power_spectrum.py
index c596d1611..2d979f4f3 100644
--- a/python/rascaline/tests/utils/power_spectrum.py
+++ b/python/rascaline/tests/utils/power_spectrum.py
@@ -65,6 +65,14 @@ def test_power_spectrum(calculator) -> None:
     assert n_props_actual == n_props_expected
 
 
+def test_forward() -> None:
+    """Test that forwar results in the same as compute."""
+    ps_compute = PowerSpectrum(soap_calculator()).compute(SystemForTests())
+    ps_forward = PowerSpectrum(soap_calculator()).forward(SystemForTests())
+
+    assert ps_compute.keys == ps_forward.keys
+
+
 def test_error_max_angular():
     """Test error raise if max_angular are different."""
     hypers_2 = HYPERS.copy()