ht.array() default to copy=None (e.g., only if necessary)

heat/core/factories.py

@@ -151,7 +151,7 @@ def arange(
 def array(
     obj: Iterable,
     dtype: Optional[Type[datatype]] = None,
-    copy: bool = True,
+    copy: bool = None,
     ndmin: int = 0,
     order: str = "C",
     split: Optional[int] = None,
@@ -172,8 +172,9 @@ def array(
         to hold the objects in the sequence. This argument can only be used to ‘upcast’ the array. For downcasting, use
         the :func:`~heat.core.dndarray.astype` method.
     copy : bool, optional
-        If ``True`` (default), then the object is copied. Otherwise, a copy will only be made if obj is a nested
-        sequence or if a copy is needed to satisfy any of the other requirements, e.g. ``dtype``.
+        If ``True``, the input object is copied.
+        If ``False``, input which supports the buffer protocol is never copied.
+        If ``None`` (default), the function reuses the existing memory buffer if possible, and copies otherwise.
     ndmin : int, optional
         Specifies the minimum number of dimensions that the resulting array should have. Ones will, if needed, be
         attached to the shape if ``ndim > 0`` and prefaced in case of ``ndim < 0`` to meet the requirement.
@@ -198,6 +199,8 @@ def array(
     ------
     NotImplementedError
         If order is one of the NumPy options ``'K'`` or ``'A'``.
+    ValueError
+        If ``copy`` is False but a copy is necessary.
 
     Examples
     --------
@@ -285,18 +288,16 @@ def array(
          [torch.LongStorage of size 6]
     """
     # array already exists; no copy
-    if (
-        isinstance(obj, DNDarray)
-        and not copy
-        and (dtype is None or dtype == obj.dtype)
-        and (split is None or split == obj.split)
-        and (is_split is None or is_split == obj.split)
-        and (device is None or device == obj.device)
-    ):
-        return obj
-
-    # extract the internal tensor in case of a heat tensor
     if isinstance(obj, DNDarray):
+        if not copy:
+            if (
+                (dtype is None or dtype == obj.dtype)
+                and (split is None or split == obj.split)
+                and (is_split is None or is_split == obj.split)
+                and (device is None or device == obj.device)
+            ):
+                return obj
+        # extract the internal tensor
         obj = obj.larray
 
     # sanitize the data type
@@ -324,13 +325,28 @@ def array(
             except RuntimeError:
                 raise TypeError("invalid data of type {}".format(type(obj)))
     else:
-        if not isinstance(obj, DNDarray):
+        if copy is False and not np.isscalar(obj) and not isinstance(obj, (Tuple, List)):
+            # Python array-API compliance, cf. https://data-apis.org/array-api/2022.12/API_specification/generated/array_api.asarray.html
+            if not (
+                (dtype is None or dtype == types.canonical_heat_type(obj.dtype))
+                and (
+                    device is None
+                    or device.torch_device
+                    == str(getattr(obj, "device", devices.get_device().torch_device))
+                )
+            ):
+                raise ValueError(
+                    "argument `copy` is set to False, but copy of input object is necessary. \n Set copy=None to reuse the memory buffer whenever possible and allow for copies otherwise."
+                )
+        try:
             obj = torch.as_tensor(
                 obj,
                 device=device.torch_device
                 if device is not None
                 else devices.get_device().torch_device,
             )
+        except RuntimeError:
+            raise TypeError("invalid data of type {}".format(type(obj)))
 
     # infer dtype from obj if not explicitly given
     if dtype is None:
@@ -377,9 +393,11 @@ def array(
 
     # content shall be split, chunk the passed data object up
     if split is not None:
+        # only keep local slice
         _, _, slices = comm.chunk(gshape, split)
-        obj = obj[slices].clone()
-        obj = sanitize_memory_layout(obj, order=order)
+        _ = obj[slices].clone()
+        del obj
+        obj = sanitize_memory_layout(_, order=order)
     # check with the neighboring rank whether the local shape would fit into a global shape
     elif is_split is not None:
         obj = sanitize_memory_layout(obj, order=order)
@@ -489,7 +507,7 @@ def asarray(
     >>> ht.asarray(a, dtype=ht.float64) is a
     False
     """
-    return array(obj, dtype=dtype, copy=False, order=order, is_split=is_split, device=device)
+    return array(obj, dtype=dtype, copy=None, order=order, is_split=is_split, device=device)
 
 
 def empty(

heat/core/manipulations.py

@@ -650,11 +650,13 @@ def diagonal(a: DNDarray, offset: int = 0, dim1: int = 0, dim2: int = 1) -> DNDa
         split = len(shape) - 1
 
     if a.split is None or a.split not in (dim1, dim2):
-        result = torch.diagonal(a.larray, offset=offset, dim1=dim1, dim2=dim2)
+        result = torch.diagonal(a.larray, offset=offset, dim1=dim1, dim2=dim2).contiguous()
     else:
         vz = 1 if a.split == dim1 else -1
         off, _, _ = a.comm.chunk(a.shape, a.split)
-        result = torch.diagonal(a.larray, offset=offset + vz * off, dim1=dim1, dim2=dim2)
+        result = torch.diagonal(
+            a.larray, offset=offset + vz * off, dim1=dim1, dim2=dim2
+        ).contiguous()
     return factories.array(result, dtype=a.dtype, is_split=split, device=a.device, comm=a.comm)
 
 

heat/core/tests/test_factories.py

@@ -307,6 +307,10 @@ def test_array(self):
         # invalid communicator
         with self.assertRaises(TypeError):
             ht.array((4,), comm={})
+        # copy=False but copy is necessary
+        data = np.arange(10)
+        with self.assertRaises(ValueError):
+            ht.array(data, dtype=ht.int32, copy=False)
 
         # data already distributed but don't match in shape
         if self.get_size() > 1:

heat/core/types.py

@@ -514,6 +514,9 @@ def canonical_heat_type(a_type: Union[str, Type[datatype], Any]) -> Type[datatyp
     except TypeError:
         pass
 
+    # extract type of numpy.dtype
+    a_type = getattr(a_type, "type", a_type)
+
     # try to look the corresponding type up
     try:
         return __type_mappings[a_type]