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

{metatensor_core-0.2.2 → metatensor_core-0.2.4}/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.4")
 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.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metatensor-core
-Version: 0.2.2
+Version: 0.2.4
 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
@@ -27,6 +27,7 @@ Description-Content-Type: text/x-rst
 License-File: LICENSE
 License-File: AUTHORS
 Requires-Dist: numpy
+Requires-Dist: ctypes-dlpack
 Dynamic: author
 Dynamic: license-file
 

metatensor_core-0.2.4/git_version_info

@@ -0,0 +1,2 @@
+0
+git.417f531

{metatensor_core-0.2.2 → metatensor_core-0.2.4}/metatensor/_block.py

@@ -48,17 +48,14 @@ class TensorBlock:
     ...     properties=Labels("properties", np.array([[0], [1], [2]])),
     ... )
     >>> block
-    TensorBlock
-        samples (2): ['samples']
-        components (): []
-        properties (3): ['properties']
-        gradients: None
+    TensorBlock with shape (2, 3)
+        samples: [samples]
+        properties: [properties]
     >>> block.samples
-    Labels(
+    Labels
         samples
            4
            2
-    )
     >>> block.values[block.samples.position([2])]
     array([3, 5, 6])
     """
@@ -254,31 +251,82 @@ class TensorBlock:
             # The block has been released
             return "TensorBlock(<empty>)"
 
+        def _format_names(names):
+            return "[" + ", ".join(names) + "]"
+
         if len(self._gradient_parameters) != 0:
-            s = f"Gradient TensorBlock ('{'/'.join(self._gradient_parameters)}')\n"
-        else:
-            s = "TensorBlock\n"
-        s += f"    samples ({len(self.samples)}): {str(list(self.samples.names))}"
-        s += "\n"
-        s += "    components ("
-        s += ", ".join([str(len(c)) for c in self.components])
-        s += "): ["
-        for ic in self.components:
-            for name in ic.names[:]:
-                s += "'" + name + "', "
-        if len(self.components) > 0:
-            s = s[:-2]
-        s += "]\n"
-        s += f"    properties ({len(self.properties)}): "
-        s += f"{str(list(self.properties.names))}\n"
-        s += "    gradients: "
-        if len(self.gradients_list()) > 0:
-            s += f"{str(list(self.gradients_list()))}"
+            s = (
+                f"TensorBlock gradient for "
+                f"'{'/'.join(self._gradient_parameters)}', "
+                f"with shape {tuple(self.values.shape)}\n"
+            )
         else:
-            s += "None"
+            s = f"TensorBlock with shape {tuple(self.values.shape)}\n"
+
+        s += f"    samples: {_format_names(self.samples.names)}\n"
+
+        component_names = []
+        for component in self.components:
+            component_names.extend(component.names)
+        if len(component_names) != 0:
+            s += f"    components: {_format_names(component_names)}\n"
+
+        s += f"    properties: {_format_names(self.properties.names)}"
+
+        gradients = self.gradients_list()
+        if len(gradients) != 0:
+            s += "\n\n    gradients:"
+            max_len = max(len(p) for p in gradients)
+            for parameter in gradients:
+                gradient = self.gradient(parameter)
+                grad_shape = tuple(int(v) for v in gradient.values.shape)
+                s += (
+                    f"\n        {parameter.ljust(max_len)} "
+                    f"=> TensorBlock with shape {grad_shape}"
+                )
 
         return s
 
+    def _repr_html_(self) -> str:
+        """HTML representation for Jupyter notebooks."""
+        if not self._ptr:
+            # The block has been released
+            return "TensorBlock(<empty>)"
+
+        from metatensor._html import TensorBlockData, block_html
+
+        def _block_data(block):
+            gradients = {}
+            for parameter in block.gradients_list():
+                grad = block.gradient(parameter)
+                gradients[parameter] = _block_data(grad)
+
+            return TensorBlockData(
+                values_shape=tuple(block.values.shape),
+                samples=(block.samples.names, block.samples.values),
+                components=[(c.names, c.values) for c in block.components],
+                properties=(block.properties.names, block.properties.values),
+                gradients=gradients,
+            )
+
+        data = _block_data(self)
+
+        body = block_html(data, module="metatensor")
+
+        class_name = "metatensor.TensorBlock"
+        if len(self._gradient_parameters) != 0:
+            gradient_for = "/".join(self._gradient_parameters)
+            rest = (
+                f" gradient for '{gradient_for}', with shape {tuple(self.values.shape)}"
+            )
+        else:
+            rest = f" with shape {tuple(self.values.shape)}"
+
+        header = f"<div><strong>{class_name}</strong>{rest}</div>"
+        hr = "<hr style='border: none; border-top: 1px solid #888; margin: 0;'>"
+
+        return header + hr + body
+
     def __eq__(self, other):
         from metatensor.operations import equal_block
 
@@ -398,19 +446,16 @@ class TensorBlock:
 
         >>> positions_gradient = block.gradient("positions")
         >>> print(positions_gradient)
-        Gradient TensorBlock ('positions')
-            samples (2): ['sample', 'atom']
-            components (3, 1): ['direction', 'component']
-            properties (5): ['property']
-            gradients: None
-
+        TensorBlock gradient for 'positions', with shape (2, 3, 1, 5)
+            samples: [sample, atom]
+            components: [direction, component]
+            properties: [property]
         >>> cell_gradient = block.gradient("cell")
         >>> print(cell_gradient)
-        Gradient TensorBlock ('cell')
-            samples (2): ['sample']
-            components (3, 3, 1): ['direction_1', 'direction_2', 'component']
-            properties (5): ['property']
-            gradients: None
+        TensorBlock gradient for 'cell', with shape (2, 3, 3, 1, 5)
+            samples: [sample]
+            components: [direction_1, direction_2, component]
+            properties: [property]
         """
         gradient_block = ctypes.POINTER(mts_block_t)()
 
@@ -463,11 +508,13 @@ class TensorBlock:
         ... )
         >>> block.add_gradient("parameter", gradient)
         >>> print(block)
-        TensorBlock
-            samples (3): ['system']
-            components (1): ['component']
-            properties (1): ['property']
-            gradients: ['parameter']
+        TensorBlock with shape (3, 1, 1)
+            samples: [system]
+            components: [component]
+            properties: [property]
+
+            gradients:
+                parameter => TensorBlock with shape (2, 1, 1)
         """
         if self._parent is not None:
             raise ValueError(

{metatensor_core-0.2.2 → metatensor_core-0.2.4}/metatensor/_c_api.py

@@ -12,41 +12,43 @@ import ctypes
 import platform
 from ctypes import CFUNCTYPE, POINTER
 
+from ctypes_dlpack import DLDataType, DLDevice, DLManagedTensorVersioned, DLPackVersion
 
-class EnumType(type(ctypes.c_int32)):
-    def __new__(metacls, name, bases, dict):
-        if "_members_" not in dict:
-            _members_ = {}
-            for key, value in dict.items():
-                if not key.startswith("_"):
-                    _members_[key] = value
 
-            dict["_members_"] = _members_
+class _EnumType(type(ctypes.c_int32)):
+    def __new__(metacls, name, bases, namespace):
+        if "_members_" not in namespace:
+            members = {}
+            for key, value in namespace.items():
+                if not key.startswith("_"):
+                    members[key] = value
+            namespace["_members_"] = members
         else:
-            _members_ = dict["_members_"]
+            members = namespace["_members_"]
 
-        dict["_reverse_map_"] = {v: k for k, v in _members_.items()}
-        cls = type(ctypes.c_int32).__new__(metacls, name, bases, dict)
-        for key, value in cls._members_.items():
-            globals()[key] = value
-        return cls
+        namespace["_reverse_map_"] = {v: k for k, v in members.items()}
+        return type(ctypes.c_int32).__new__(metacls, name, bases, namespace)
 
     def __repr__(self):
-        return "<Enum %s>" % self.__name__
+        return f"<Enum {self.__name__}>"
 
 
-class Enum(ctypes.c_int32, metaclass=EnumType):
+class _Enum(ctypes.c_int32, metaclass=_EnumType):
     _members_ = {}
 
     def __repr__(self):
-        value_str = self._reverse_map_.get(self.value, str(self.value))
-        return f"{self.__class__.__name__}.{value_str}"
+        value_name = self._reverse_map_.get(self.value, str(self.value))
+        return f"{self.__class__.__name__}.{value_name}"
 
     def __eq__(self, other):
         if isinstance(other, int):
             return self.value == other
+        if type(self) is type(other):
+            return self.value == other.value
+        return NotImplemented
 
-        return type(self) is type(other) and self.value == other.value
+    def __hash__(self):
+        return hash(self.value)
 
 
 arch = platform.architecture()[0]
@@ -57,47 +59,7 @@ elif arch == "64bit":
 
 
 
-class DLDeviceType(Enum):
-    kDLCPU = 1
-    kDLCUDA = 2
-    kDLCUDAHost = 3
-    kDLOpenCL = 4
-    kDLVulkan = 7
-    kDLMetal = 8
-    kDLVPI = 9
-    kDLROCM = 10
-    kDLROCMHost = 11
-    kDLExtDev = 12
-    kDLCUDAManaged = 13
-    kDLOneAPI = 14
-    kDLWebGPU = 15
-    kDLHexagon = 16
-    kDLMAIA = 17
-    kDLTrn = 18
-
-
-class DLDataTypeCode(Enum):
-    kDLInt = 0
-    kDLUInt = 1
-    kDLFloat = 2
-    kDLOpaqueHandle = 3
-    kDLBfloat = 4
-    kDLComplex = 5
-    kDLBool = 6
-    kDLFloat8_e3m4 = 7
-    kDLFloat8_e4m3 = 8
-    kDLFloat8_e4m3b11fnuz = 9
-    kDLFloat8_e4m3fn = 10
-    kDLFloat8_e4m3fnuz = 11
-    kDLFloat8_e5m2 = 12
-    kDLFloat8_e5m2fnuz = 13
-    kDLFloat8_e8m0fnu = 14
-    kDLFloat6_e2m3fn = 15
-    kDLFloat6_e3m2fn = 16
-    kDLFloat4_e2m1fn = 17
-
-
-class mts_status_t(Enum):
+class mts_status_t(_Enum):
     MTS_SUCCESS = 0
     MTS_INVALID_PARAMETER_ERROR = 1
     MTS_IO_ERROR = 2
@@ -107,30 +69,6 @@ class mts_status_t(Enum):
     MTS_INTERNAL_ERROR = 255
 
 
-class DLPackVersion(ctypes.Structure):
-    pass
-
-
-class DLDevice(ctypes.Structure):
-    pass
-
-
-class DLDataType(ctypes.Structure):
-    pass
-
-
-class DLTensor(ctypes.Structure):
-    pass
-
-
-class DLManagedTensor(ctypes.Structure):
-    pass
-
-
-class DLManagedTensorVersioned(ctypes.Structure):
-    pass
-
-
 class mts_block_t(ctypes.Structure):
     pass
 
@@ -151,56 +89,11 @@ class mts_array_t(ctypes.Structure):
     pass
 
 
-DLPackManagedTensorAllocator = CFUNCTYPE(ctypes.c_int, POINTER(DLTensor), POINTER(POINTER(DLManagedTensorVersioned)), ctypes.c_void_p, CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p))
-DLPackManagedTensorFromPyObjectNoSync = CFUNCTYPE(ctypes.c_int, ctypes.c_void_p, POINTER(POINTER(DLManagedTensorVersioned)))
-DLPackDLTensorFromPyObjectNoSync = CFUNCTYPE(ctypes.c_int, ctypes.c_void_p, POINTER(DLTensor))
-DLPackCurrentWorkStream = CFUNCTYPE(ctypes.c_int, DLDeviceType, ctypes.c_int32, POINTER(POINTER(None)))
-DLPackManagedTensorToPyObjectNoSync = CFUNCTYPE(ctypes.c_int, POINTER(DLManagedTensorVersioned), POINTER(POINTER(None)))
 mts_data_origin_t = ctypes.c_uint64
 mts_realloc_buffer_t = CFUNCTYPE(ctypes.c_char_p, ctypes.c_void_p, ctypes.c_char_p, c_uintptr_t)
 mts_create_array_callback_t = CFUNCTYPE(mts_status_t, POINTER(c_uintptr_t), c_uintptr_t, DLDataType, POINTER(mts_array_t))
 
 
-DLPackVersion._fields_ = [
-    ("major", ctypes.c_uint32),
-    ("minor", ctypes.c_uint32),
-]
-
-DLDevice._fields_ = [
-    ("device_type", DLDeviceType),
-    ("device_id", ctypes.c_int32),
-]
-
-DLDataType._fields_ = [
-    ("code", ctypes.c_uint8),
-    ("bits", ctypes.c_uint8),
-    ("lanes", ctypes.c_uint16),
-]
-
-DLTensor._fields_ = [
-    ("data", ctypes.c_void_p),
-    ("device", DLDevice),
-    ("ndim", ctypes.c_int32),
-    ("dtype", DLDataType),
-    ("shape", POINTER(ctypes.c_int64)),
-    ("strides", POINTER(ctypes.c_int64)),
-    ("byte_offset", ctypes.c_uint64),
-]
-
-DLManagedTensor._fields_ = [
-    ("dl_tensor", DLTensor),
-    ("manager_ctx", ctypes.c_void_p),
-    ("deleter", CFUNCTYPE(None, POINTER(DLManagedTensor))),
-]
-
-DLManagedTensorVersioned._fields_ = [
-    ("version", DLPackVersion),
-    ("manager_ctx", ctypes.c_void_p),
-    ("deleter", CFUNCTYPE(None, POINTER(DLManagedTensorVersioned))),
-    ("flags", ctypes.c_uint64),
-    ("dl_tensor", DLTensor),
-]
-
 mts_data_movement_t._fields_ = [
     ("sample_in", c_uintptr_t),
     ("sample_out", c_uintptr_t),

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

@@ -2,29 +2,23 @@ import ctypes
 from typing import NewType, Union
 
 import numpy as np
-
-from .._c_api import (
+from ctypes_dlpack import (
     DLDataType,
     DLDevice,
     DLDeviceType,
     DLManagedTensorVersioned,
+    DLPackArray,
     DLPackVersion,
+    array_as_dlpack,
+)
+
+from .._c_api import (
     c_uintptr_t,
     mts_array_t,
     mts_data_movement_t,
     mts_data_origin_t,
 )
 from .._status import catch_exceptions, check_status
-from ._dlpack import (
-    DLPACK_NAME,
-    DLPACK_VERSIONED_NAME,
-    PYTHON_API,
-    USED_DLPACK_NAME,
-    USED_DLPACK_VERSIONED_NAME,
-    DLManagedTensor,
-    DLPackArray,
-    wrap_unversioned_as_versioned,
-)
 
 
 try:
@@ -37,9 +31,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`.
@@ -500,107 +494,21 @@ def _mts_array_move_data(
 @catch_exceptions
 def _mts_array_as_dlpack(this, dl_managed_tensor_ptr_ptr, device, stream, max_version):
     """
-    Implementation of `mts_array_t.as_dlpack`.
+    Implementation of ``mts_array_t.as_dlpack``.
 
-    This function calls the array's __dlpack__ method, gets the PyCapsule,
-    extracts the raw pointer, and transfers
-    ownership to the C-API caller.
-
-    When the PyCapsule contains a `dltensor_versioned`, we just need to rename
-    the capsule to `used_dltensor_versioned` after giving the data to the C API.
-
-    When the PyCapsule contains a deprecated `dltensor`, we re-wrap the data
-    inside `DLManagedTensorVersioned`.
+    This function calls the array's __dlpack__ method, gets the PyCapsule, extracts the
+    raw pointer, and transfers ownership to the C-API caller.
     """
     wrapper = _KNOWN_ARRAY_WRAPPERS[this]
-    array = wrapper.array
-
-    dl_device = (device.device_type.value, device.device_id)
-    stream = stream.contents if stream else None
-    max_version = (max_version.major, max_version.minor)
 
-    capsule = None
-
-    try:
-        # Try requesting versioned DLPack
-        capsule = array.__dlpack__(
-            stream=stream, max_version=max_version, dl_device=dl_device
-        )
-    except Exception as _:
-        # Fallback to legacy signatures. Each fallback drops one parameter.
-        try:
-            capsule = array.__dlpack__(stream=stream, dl_device=dl_device)
-        except Exception:
-            try:
-                capsule = array.__dlpack__(dl_device=dl_device)
-            except Exception:
-                capsule = array.__dlpack__()
-
-    capsule_name = PYTHON_API.PyCapsule_GetName(capsule)
-
-    # Versioned Capsule
-    if capsule_name == DLPACK_VERSIONED_NAME:
-        pointer = PYTHON_API.PyCapsule_GetPointer(capsule, DLPACK_VERSIONED_NAME)
-        if not pointer:
-            raise
-
-        versioned_ptr = ctypes.cast(pointer, ctypes.POINTER(DLManagedTensorVersioned))
-        actual_device = versioned_ptr.contents.dl_tensor.device
-        if (
-            actual_device.device_type != device.device_type
-            or actual_device.device_id != device.device_id
-        ):
-            raise ValueError(
-                f"DLPack device mismatch: expected type={device.device_type} "
-                f"id={device.device_id}, got type={actual_device.device_type} "
-                f"id={actual_device.device_id}"
-            )
-
-        status = PYTHON_API.PyCapsule_SetName(capsule, USED_DLPACK_VERSIONED_NAME)
-        if status != 0:
-            raise
-
-        dl_managed_tensor_ptr_ptr[0] = versioned_ptr
-        return
-
-    # Legacy Capsule
-    if capsule_name == DLPACK_NAME:
-        pointer = PYTHON_API.PyCapsule_GetPointer(capsule, DLPACK_NAME)
-        if not pointer:
-            raise
-
-        unversioned_ptr = ctypes.cast(pointer, ctypes.POINTER(DLManagedTensor))
-
-        actual_device = unversioned_ptr.contents.dl_tensor.device
-        if (
-            actual_device.device_type != device.device_type
-            or actual_device.device_id != device.device_id
-        ):
-            raise ValueError(
-                f"DLPack device mismatch: expected type {device.device_type}"
-                f" id {device.device_id}, "
-                f"got type {actual_device.device_type}"
-                f" id {actual_device.device_id}"
-            )
-
-        # We must do this to ensure that the deletion of the PyCapsule does not run the
-        # deleter and cause double-free issues. Instead we call legacy deleter
-        # explicitly as part of the versioned tensor's deleter.
-        status = PYTHON_API.PyCapsule_SetName(capsule, USED_DLPACK_NAME)
-        if status != 0:
-            raise
-
-        # Wrap it
-        versioned_ptr = wrap_unversioned_as_versioned(unversioned_ptr)
-        dl_managed_tensor_ptr_ptr[0] = versioned_ptr
-        return
-
-    raise ValueError(
-        "Unexpected DLPack capsule name:"
-        f" '{capsule_name.decode() if capsule_name else 'NULL'}'. "
-        f"Expected '{DLPACK_VERSIONED_NAME.decode()}' or '{DLPACK_NAME.decode()}'"
+    stream = None if not stream else stream.contents
+    versioned_ptr = array_as_dlpack(
+        wrapper.array, dl_device=device, stream=stream, max_version=max_version
     )
 
+    dl_managed_tensor_ptr_ptr[0] = versioned_ptr
+    return
+
 
 @catch_exceptions
 def _mts_array_from_dlpack(this, dl_tensor, new_array):

{metatensor_core-0.2.2 → metatensor_core-0.2.4}/metatensor/_data/_extract.py

@@ -2,12 +2,15 @@ import ctypes
 from typing import Any
 
 import numpy as np
-
-from .._c_api import (
+from ctypes_dlpack import (
     DLDevice,
     DLDeviceType,
     DLManagedTensorVersioned,
+    DLPackArray,
     DLPackVersion,
+)
+
+from .._c_api import (
     c_uintptr_t,
     mts_array_t,
     mts_data_origin_t,
@@ -20,7 +23,6 @@ from ._array import (
     _origin_pytorch,
     _register_origin,
 )
-from ._dlpack import DLPackArray, wrap_versioned_as_unversioned
 
 
 _ADDITIONAL_ORIGINS = {}
@@ -235,13 +237,7 @@ class ExternalCudaArray:
                 pass
 
         if tensor is None:
-            try:
-                tensor = torch.from_dlpack(dlpack_array)
-            except RuntimeError:
-                # Older PyTorch (< 2.4) doesn't understand versioned
-                # DLPack capsules. Fall back to an unversioned capsule.
-                unversioned = wrap_versioned_as_unversioned(dl_managed_ptr)
-                tensor = torch.from_dlpack(DLPackArray(unversioned))
+            tensor = torch.from_dlpack(dlpack_array)
 
         # keep a reference to the parent object to prevent it from being
         # garbage-collected while the tensor is alive