Add testing utilities to perform finite difference for operations

- Add finite difference test for add operation
- Add dispatch functions to add tests for torch array backend

Co-authored-by: Divya Suman <[email protected]>

python/metatensor-operations/metatensor/operations/__init__.py

@@ -60,3 +60,5 @@
 from .subtract import subtract # noqa
 from .unique_metadata import unique_metadata, unique_metadata_block # noqa
 from .zeros_like import zeros_like, zeros_like_block # noqa
+
+from . import _testing # noqa

python/metatensor-operations/metatensor/operations/_dispatch.py

@@ -46,6 +46,20 @@ def _check_all_np_ndarray(arrays):
  )
 
 
+def sum(array, axis: Optional[int] = None):
+ """
+ Returns the sum of the elements in the array at the axis.
+
+ It is equivalent of np.sum(array, axis=axis) and torch.sum(tensor, dim=axis)
+ """
+ if isinstance(array, TorchTensor):
+ return torch.sum(array, dim=axis)
+ elif isinstance(array, np.ndarray):
+ return np.sum(array, axis=axis).astype(array.dtype)
+ else:
+ raise TypeError(UNKNOWN_ARRAY_TYPE)
+
+
 def abs(array):
  """
  Returns the absolute value of the elements in the array.

python/metatensor-operations/tests/add.py

@@ -3,6 +3,30 @@
 
 import metatensor
 from metatensor import Labels, TensorBlock, TensorMap
+from metatensor.operations._testing import (
+ cartesian_cubic,
+ cartesian_linear,
+ finite_differences,
+)
+
+
+try:
+ import torch # noqa
+
+ HAS_TORCH = True
+except ImportError:
+ HAS_TORCH = False
+
+
+@pytest.fixture(scope="module", autouse=True)
+def set_random_generator():
+ """Set the random generator to same seed before each test is run.
+ Otherwise test behaviour is dependend on the order of the tests
+ in this file and the number of parameters of the test.
+ """
+ np.random.seed(1225787)
+ if HAS_TORCH:
+ torch.manual_seed(1225787)
 
 
 @pytest.fixture
@@ -258,3 +282,24 @@ def test_self_add_error():
  )
  with pytest.raises(TypeError, match=message):
  metatensor.add(tensor, np.ones((3, 4)))
+
+
+def test_add_finite_difference():
+ def add_callable(cartesian_vectors, compute_grad=False):
+ tensor1 = cartesian_linear(cartesian_vectors, compute_grad)
+ tensor2 = cartesian_cubic(cartesian_vectors, compute_grad)
+ return metatensor.add(tensor1, tensor2)
+
+ input_array = np.random.rand(5, 3)
+ finite_differences(add_callable, input_array, "positions")
+
+
+@pytest.mark.skipif(not HAS_TORCH, reason="requires torch")
+def test_torch_add_finite_difference():
+ def add_callable(cartesian_vectors, compute_grad=False):
+ tensor1 = cartesian_linear(cartesian_vectors, compute_grad)
+ tensor2 = cartesian_cubic(cartesian_vectors, compute_grad)
+ return metatensor.add(tensor1, tensor2)
+
+ input_array = torch.rand(5, 3, dtype=torch.float64)
+ finite_differences(add_callable, input_array, "positions")