Features/176 flatten

CHANGELOG.md

@@ -5,6 +5,7 @@
 - [#483](https://github.com/helmholtz-analytics/heat/pull/483) Bugfix: Underlying torch tensor moves to the right device on array initialisation
 - [#483](https://github.com/helmholtz-analytics/heat/pull/483) Bugfix:DNDarray.cpu() changes heat device to cpu
 - Update documentation theme to "Read the Docs"
+- [#501](https://github.com/helmholtz-analytics/heat/pull/501) New Feature: flatten
 
 # v0.3.0
 

heat/core/manipulations.py

@@ -14,6 +14,7 @@
     "diag",
     "diagonal",
     "expand_dims",
+    "flatten",
     "hstack",
     "resplit",
     "sort",
@@ -556,6 +557,46 @@ def expand_dims(a, axis):
     )
 
 
+def flatten(a):
+    """
+    Flattens an array into one dimension
+
+    Parameters
+    ----------
+    a : DNDarray
+        array to collapse
+    Returns
+    -------
+    ret : DNDarray
+        flattened copy
+    Examples
+    --------
+    >>> a = ht.array([[[1,2],[3,4]],[[5,6],[7,8]]])
+    >>> ht.flatten(a)
+    tensor([1,2,3,4,5,6,7,8])
+    """
+    if a.split is None:
+        return factories.array(
+            torch.flatten(a._DNDarray__array),
+            dtype=a.dtype,
+            is_split=None,
+            device=a.device,
+            comm=a.comm,
+        )
+
+    a = resplit(a, 0)
+    a = factories.array(
+        torch.flatten(a._DNDarray__array),
+        dtype=a.dtype,
+        is_split=a.split,
+        device=a.device,
+        comm=a.comm,
+    )
+    a.balance_()
+
+    return a
+
+
 def hstack(tup):
     """
     Stack arrays in sequence horizontally (column wise).

heat/core/tests/test_manipulations.py

@@ -763,6 +763,31 @@ def test_expand_dims(self):
         with self.assertRaises(ValueError):
             ht.empty((3, 4, 5), device=ht_device).expand_dims(-5)
 
+    def test_flatten(self):
+        a = ht.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], device=ht_device)
+        res = ht.array([1, 2, 3, 4, 5, 6, 7, 8], device=ht_device)
+        self.assertTrue(ht.equal(ht.flatten(a), res))
+
+        a = ht.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], split=0, device=ht_device, dtype=ht.int8)
+        res = ht.array([1, 2, 3, 4, 5, 6, 7, 8], split=0, device=ht_device, dtype=ht.int8)
+        self.assertTrue(ht.equal(ht.flatten(a), res))
+
+        a = ht.array(
+            [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], split=1, device=ht_device
+        )
+        res = ht.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], split=0, device=ht_device)
+        self.assertTrue(ht.equal(ht.flatten(a), res))
+
+        a = ht.array(
+            [[[False, False], [False, True]], [[True, False], [True, True]]],
+            split=2,
+            device=ht_device,
+        )
+        res = ht.array(
+            [False, False, False, True, True, False, True, True], split=0, device=ht_device
+        )
+        self.assertTrue(ht.equal(ht.flatten(a), res))
+
     def test_hstack(self):
         # cases to test:
         # MM===================================