clebsch gordan submodule - implementation of TorchScript interface

python/rascaline-torch/tests/utils/correlate_density.py

@@ -16,6 +16,7 @@
 DATA_ROOT = os.path.join(os.path.dirname(__file__), "data")
 
 
+@torch.jit.script
 def is_tensor_map(obj: Any):
     return isinstance(obj, TensorMap)
 
@@ -38,7 +39,16 @@ def is_tensor_map(obj: Any):
 
 
 def h2o_isolated():
-    return ase.io.read(os.path.join(DATA_ROOT, "h2o_isolated.xyz"), ":")
+    return [
+        ase.Atoms(
+            symbols=["O", "H", "H"],
+            positions=[
+                [2.56633400, 2.50000000, 2.50370100],
+                [1.97361700, 1.73067300, 2.47063400],
+                [1.97361700, 3.26932700, 2.47063400],
+            ],
+        )
+    ]
 
 
 def spherical_expansion(frames: List[ase.Atoms]):
@@ -93,7 +103,6 @@ def test_jit_save_load():
         torch.jit.save(scripted_correlate_density, buffer)
         buffer.seek(0)
         torch.jit.load(buffer)
-        buffer.close()
 
 
 def test_save_load():
@@ -110,4 +119,3 @@ def test_save_load():
         torch.save(corr_calculator, buffer)
         buffer.seek(0)
         torch.load(buffer)
-        buffer.close()

python/rascaline/rascaline/utils/_dispatch.py

@@ -32,15 +32,6 @@ def _check_all_torch_tensor(arrays: List[TorchTensor]):
             )
 
 
-# def _check_all_torch_tensor(arrays: List[TorchTensor]):
-#     for array in arrays:
-#         if not isinstance(array, TorchTensor):
-#             raise TypeError(
-#                 f"expected argument to be a torch.Tensor, but got
-#                 {type(array)}"
-#             )
-
-
 def _check_all_np_ndarray(arrays):
     for array in arrays:
         if not isinstance(array, np.ndarray):
@@ -49,14 +40,6 @@ def _check_all_np_ndarray(arrays):
             )
 
 
-# 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 concatenate(arrays: List[TorchTensor], axis: int):
     """
     Concatenate a group of arrays along a given axis.
@@ -78,16 +61,6 @@ def concatenate(arrays: List[TorchTensor], axis: int):
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
-# def concatenate(arrays, axis: Optional[int] = 0):
-#     """Concatenate arrays along an axis."""
-#     if isinstance(arrays[0], TorchTensor):
-#         return torch.cat(arrays, dim=axis)
-#     elif isinstance(arrays[0], np.ndarray):
-#         return np.concatenate(arrays, axis=axis)
-#     else:
-#         raise TypeError(UNKNOWN_ARRAY_TYPE)
-
-
 def empty_like(array, shape: Optional[List[int]] = None, requires_grad: bool = False):
     """
     Create an uninitialized array, with the given ``shape``, and similar dtype,
@@ -114,33 +87,6 @@ def empty_like(array, shape: Optional[List[int]] = None, requires_grad: bool = F
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
-# def empty_like(array, shape: Optional[List[int]] = None, requires_grad: bool = False):
-#     """
-#     Create an empty array, 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.empty_like(array, shape=shape)``.
-#     """
-#     if isinstance(array, TorchTensor):
-#         if shape is None:
-#             shape = array.size()
-
-#         return torch.empty(
-#             shape,
-#             dtype=array.dtype,
-#             layout=array.layout,
-#             device=array.device,
-#         ).requires_grad_(requires_grad)
-#     elif isinstance(array, np.ndarray):
-#         return np.empty_like(array, shape=shape, subok=False)
-#     else:
-#         raise TypeError(UNKNOWN_ARRAY_TYPE)
-
-
 def list_to_array(array, data: List[List[int]]):
     """Create an object from data with the same type as ``array``."""
     if isinstance(array, TorchTensor):
@@ -195,14 +141,6 @@ def unique(array, axis: Optional[int] = None):
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
-# def unique(array, axis: Optional[int] = None):
-#     """Find the unique elements of an array."""
-#     if isinstance(array, TorchTensor):
-#         return torch.unique(array, dim=axis)
-#     elif isinstance(array, np.ndarray):
-#         return np.unique(array, axis=axis)
-
-
 def zeros_like(array, shape: Optional[List[int]] = None, requires_grad: bool = False):
     """
     Create an array filled with zeros, with the given ``shape``, and similar
@@ -232,33 +170,6 @@ def zeros_like(array, shape: Optional[List[int]] = None, requires_grad: bool = F
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
-# def zeros_like(array, shape: Optional[List[int]] = None, requires_grad: bool = 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)
-
-
 def where(array):
     """Return the indices where `array` is True.