metatensor-core 0.2.2__tar.gz → 0.2.3__tar.gz

{metatensor_core-0.2.2 → metatensor_core-0.2.3}/CMakeLists.txt

@@ -14,7 +14,7 @@ set(METATENSOR_CORE_SOURCE_DIR "" CACHE PATH "Path to the sources of metatensor-
 
 file(REMOVE ${CMAKE_INSTALL_PREFIX}/_external.py)
 
-set(REQUIRED_METATENSOR_VERSION "0.2.2")
+set(REQUIRED_METATENSOR_VERSION "0.2.3")
 if(${METATENSOR_CORE_PYTHON_USE_EXTERNAL_LIB})
     # when building a source checkout, update version to include git information
     # this will not apply when building a sdist

{metatensor_core-0.2.2 → metatensor_core-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metatensor-core
-Version: 0.2.2
+Version: 0.2.3
 Summary: Python bindings for metatensor
 Author: Guillaume Fraux, Rohit Goswami, Philip Loche, Filippo Bigi, Joseph W. Abbott, Davide Tisi, Alexander Goscinski, Michele Ceriotti
 License-Expression: BSD-3-Clause

metatensor_core-0.2.3/git_version_info

@@ -0,0 +1,2 @@
+0
+git.604d7b2

{metatensor_core-0.2.2 → metatensor_core-0.2.3}/metatensor/_data/_array.py

@@ -37,9 +37,9 @@ except ImportError:
 if HAS_TORCH:
     # This NewType is only used for typechecking and documentation purposes. If you are
     # trying to add support for new array types, see `data.array.ArrayWrapper` instead.
-    Array = NewType("Array", Union[np.ndarray, torch.Tensor])
+    Array = NewType("Array", Union[np.ndarray, torch.Tensor, mts_array_t])
 else:
-    Array = NewType("Array", np.ndarray)
+    Array = NewType("Array", Union[np.ndarray, mts_array_t])
 
 Array.__doc__ = """
 An ``Array`` contains the actual data stored in a :py:class:`metatensor.TensorBlock`.

{metatensor_core-0.2.2 → metatensor_core-0.2.3}/metatensor/_data/_dlpack.py

@@ -264,16 +264,6 @@ class DLPackArray:
         if stream is not None:
             raise RuntimeError("only `stream=None` is supported")
 
-        if self._versioned:
-            version = pointer[0].version
-            version = (version.major, version.minor)
-            if max_version is not None:
-                if version[0] > max_version[0]:
-                    raise RuntimeError(
-                        f"requested DLPack version {max_version}, but tensor has "
-                        f"version {version}"
-                    )
-
         if dl_device is not None and dl_device != self.__dlpack_device__():
             raise RuntimeError("device conversion is not supported")
 
@@ -282,7 +272,22 @@ class DLPackArray:
 
         self._pointer = None
         if self._versioned:
-            capsule = make_dlpack_versioned_capsule(pointer)
+            version = pointer[0].version
+            version_tuple = (version.major, version.minor)
+
+            if max_version is not None and version_tuple[0] > max_version[0]:
+                raise RuntimeError(
+                    f"requested DLPack version {max_version}, but tensor has "
+                    f"version {version_tuple}"
+                )
+
+            if max_version is not None and max_version >= (1, 0):
+                capsule = make_dlpack_versioned_capsule(pointer)
+            else:
+                # Consumer doesn't signal DLPack v1 support (e.g. PyTorch < 2.9.0),
+                # wrap as unversioned for compatibility
+                unversioned_ptr = wrap_versioned_as_unversioned(pointer)
+                capsule = make_dlpack_unversioned_capsule(unversioned_ptr)
         else:
             capsule = make_dlpack_unversioned_capsule(pointer)
 

{metatensor_core-0.2.2 → metatensor_core-0.2.3}/metatensor/_labels.py

@@ -8,13 +8,15 @@ import numpy as np
 
 from . import _data
 from ._c_api import (
+    DLDataType,
+    DLDataTypeCode,
     c_uintptr_t,
     mts_array_t,
     mts_labels_t,
 )
 from ._c_lib import _get_library
 from ._data import Array, Device
-from ._status import check_pointer
+from ._status import check_pointer, check_status
 
 
 try:
@@ -138,6 +140,72 @@ class LabelsEntry:
         return not self.__eq__(other)
 
 
+def _labels_values_mts_array(array: Array, size: int) -> (mts_array_t, Array):
+    if isinstance(array, mts_array_t):
+        shape_ptr = ctypes.POINTER(c_uintptr_t)()
+        shape_count = c_uintptr_t()
+        check_status(
+            array.shape(array.ptr, ctypes.byref(shape_ptr), ctypes.byref(shape_count))
+        )
+        shape = [shape_ptr[i] for i in range(shape_count.value)]
+
+        if len(shape) != 2:
+            raise ValueError("`values` must be a 2D array")
+
+        if shape[1] != size:
+            raise ValueError(
+                "`names` must have an entry for each column of the `values` array"
+            )
+
+        dtype = DLDataType()
+        check_status(array.dtype(array.ptr, ctypes.byref(dtype)))
+        if dtype.code != DLDataTypeCode.kDLInt or dtype.bits != 32 or dtype.lanes != 1:
+            raise TypeError("Labels values must be of dtype int32")
+
+        return array, None
+    else:
+        values = array
+        backend = _data.array_backend(values)
+
+        if len(values) == 0:
+            # Explicitly define empty labels
+            values = np.empty((0, size), dtype=np.int32)
+            values = _data.array_change_backend(values, backend)
+
+        if len(values.shape) != 2:
+            # make sure the array is 2D
+            raise ValueError("`values` must be a 2D array")
+
+        if values.shape[1] != size:
+            raise ValueError(
+                "`names` must have an entry for each column of the `values` array"
+            )
+
+        if backend == "numpy":
+            try:
+                # We need to make sure the data has the right type
+                values = values.astype(
+                    np.int32,
+                    casting="same_kind",
+                    subok=False,
+                    copy=False,
+                )
+                # values should not be writeable
+                values.flags.writeable = False
+            except TypeError as e:
+                raise TypeError("Labels values must be convertible to integers") from e
+        elif backend == "torch":
+            if values.requires_grad:
+                raise ValueError("Labels values can not require gradients")
+
+            if not _data.array_dtype_is_integer(values):
+                raise TypeError("Labels values must be of integer dtype")
+
+            values = values.to(dtype=torch.int32, copy=False)
+
+        return _data.create_mts_array(values), values
+
+
 class Labels:
     """
     A set of labels carrying metadata associated with a :py:class:`TensorMap`.
@@ -213,7 +281,7 @@ class Labels:
     def __init__(
         self,
         names: Union[str, Sequence[str]],
-        values: np.ndarray,
+        values: Array,
         assume_unique: bool = False,
     ):
         """
@@ -227,59 +295,20 @@ class Labels:
             either by construction or because you checked.
         """
 
-        names = _normalize_names_type(names)
-
         if not isinstance(values, Array.__supertype__):
-            raise ValueError("`values` must be a numpy ndarray or torch tensor")
-
-        backend = _data.array_backend(values)
-
-        if len(values) == 0:
-            # Explicitly define empty labels
-            values = np.empty((0, len(names)), dtype=np.int32)
-            values = _data.array_change_backend(values, backend)
-
-        if len(values.shape) != 2:
-            # make sure the array is 2D
-            raise ValueError("`values` must be a 2D array")
-
-        if len(names) != values.shape[1]:
             raise ValueError(
-                "`names` must have an entry for each column of the `values` array"
+                "`values` must be a numpy ndarray, torch tensor, or mts_array_t, "
+                f"got {type(values)}"
             )
 
-        if backend == "numpy":
-            try:
-                # We need to make sure the data has the right type
-                values = values.astype(
-                    np.int32,
-                    casting="same_kind",
-                    subok=False,
-                    copy=False,
-                )
-                # values should not be writeable
-                values.flags.writeable = False
-            except TypeError as e:
-                raise TypeError("Labels values must be convertible to integers") from e
-        elif backend == "torch":
-            if values.requires_grad:
-                raise ValueError("Labels values can not require gradients")
-
-            if not _data.array_dtype_is_integer(values):
-                raise TypeError("Labels values must be of integer dtype")
-
-            values = values.to(dtype=torch.int32, copy=False)
-
-        if values.ndim == 1:
-            values = values.reshape(-1, len(names))
-
+        names = _normalize_names_type(names)
         self._lib = _get_library()
+        array, values = _labels_values_mts_array(values, len(names))
+
         c_names = ctypes.ARRAY(ctypes.c_char_p, len(names))()
         for i, n in enumerate(names):
             c_names[i] = n.encode("utf8")
 
-        array = _data.create_mts_array(values)
-
         if assume_unique:
             ptr = self._lib.mts_labels_assume_unique(c_names, len(names), array)
         else:

{metatensor_core-0.2.2 → metatensor_core-0.2.3}/metatensor/_tensor.py

@@ -7,11 +7,20 @@ import numpy as np
 
 from . import _data
 from ._block import TensorBlock
-from ._c_api import c_uintptr_t, mts_array_t, mts_block_t, mts_tensormap_t
+from ._c_api import (
+    DLDevice,
+    DLDeviceType,
+    DLManagedTensorVersioned,
+    DLPackVersion,
+    c_uintptr_t,
+    mts_array_t,
+    mts_block_t,
+    mts_tensormap_t,
+)
 from ._c_lib import _get_library
 from ._data import Device, DType
 from ._labels import Labels, LabelsEntry
-from ._status import check_pointer
+from ._status import check_pointer, check_status
 
 
 class TensorMap:
@@ -417,7 +426,7 @@ class TensorMap:
         elif isinstance(selection, int):
             return self.block_by_id(selection)
         else:
-            selection = _normalize_selection(selection)
+            selection = _normalize_selection(selection, like=self.keys)
 
         keys = self.keys
         matching = keys.select(selection)
@@ -471,7 +480,7 @@ class TensorMap:
         elif isinstance(selection, int):
             return [self.block_by_id(selection)]
         else:
-            selection = _normalize_selection(selection)
+            selection = _normalize_selection(selection, like=self.keys)
 
         keys = self.keys
         matching = keys.select(selection)
@@ -833,8 +842,69 @@ def _can_cast_to_numpy_int(value):
     return np.can_cast(value, np.int32, casting="same_kind")
 
 
+def _array_like(values: List[int], like: Labels) -> mts_array_t:
+    """
+    Convert a list of integers to an array with the same dtype, device and array backend
+    as ``like``.
+
+    :param values: list of integers to convert
+    :param like: Labels whose array backend and device to match
+    :return: a new ``mts_array_t`` with the same dtype, device and array backend as
+        ``like``
+    """
+    np_values = np.array(values, dtype=np.int32).reshape(1, -1)
+    cpu_array = _data.create_mts_array(np_values)
+
+    try:
+        like_array = like._raw_values
+
+        dl_managed_ptr = ctypes.POINTER(DLManagedTensorVersioned)()
+        device = DLDevice(device_type=DLDeviceType.kDLCPU, device_id=0)
+        version = DLPackVersion(major=1, minor=0)
+        status = cpu_array.as_dlpack(
+            cpu_array.ptr,
+            ctypes.byref(dl_managed_ptr),
+            device,
+            None,
+            version,
+        )
+        check_status(status)
+
+        result_array = mts_array_t()
+        status = like_array.from_dlpack(
+            like_array.ptr,
+            dl_managed_ptr,
+            ctypes.byref(result_array),
+        )
+        check_status(status)
+
+        target_device = DLDevice()
+        status = like_array.device(like_array.ptr, ctypes.byref(target_device))
+        check_status(status)
+
+        if target_device.device_type != DLDeviceType.kDLCPU:
+            device_array = mts_array_t()
+            status = result_array.copy(
+                result_array.ptr,
+                target_device,
+                ctypes.byref(device_array),
+            )
+            check_status(status)
+
+            if result_array.destroy:
+                result_array.destroy(result_array.ptr)
+
+            result_array = device_array
+
+        return result_array
+    finally:
+        if cpu_array.destroy:
+            cpu_array.destroy(cpu_array.ptr)
+
+
 def _normalize_selection(
     selection: Union[Labels, LabelsEntry, Dict[str, int]],
+    like: Labels,
 ) -> Labels:
     if isinstance(selection, dict):
         for key, value in selection.items():
@@ -848,11 +918,11 @@ def _normalize_selection(
                 )
 
         if len(selection) == 0:
-            return Labels([], np.empty((0, 0), dtype=np.int32))
+            return Labels([], _array_like([], like))
         else:
             return Labels(
                 list(selection.keys()),
-                np.array([[np.int32(v) for v in selection.values()]], dtype=np.int32),
+                _array_like([np.int32(v) for v in selection.values()], like),
             )
 
     elif isinstance(selection, Labels):

{metatensor_core-0.2.2 → metatensor_core-0.2.3}/metatensor_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metatensor-core
-Version: 0.2.2
+Version: 0.2.3
 Summary: Python bindings for metatensor
 Author: Guillaume Fraux, Rohit Goswami, Philip Loche, Filippo Bigi, Joseph W. Abbott, Davide Tisi, Alexander Goscinski, Michele Ceriotti
 License-Expression: BSD-3-Clause

{metatensor_core-0.2.2 → metatensor_core-0.2.3}/metatensor_core.egg-info/SOURCES.txt

@@ -4,7 +4,7 @@ LICENSE
 MANIFEST.in
 README.rst
 git_version_info
-metatensor-core-cxx-0.2.2.tar.gz
+metatensor-core-cxx-0.2.3.tar.gz
 pyproject.toml
 setup.py
 metatensor/__init__.py

metatensor_core-0.2.2/git_version_info

@@ -1,2 +0,0 @@
-0
-git.cf11e6d