Features/176 flatten

CHANGELOG.md

@@ -10,6 +10,7 @@
 - [#496](https://github.com/helmholtz-analytics/heat/pull/496) New feature: flipud()
 - [#498](https://github.com/helmholtz-analytics/heat/pull/498) Feature: flip()
 - [#499](https://github.com/helmholtz-analytics/heat/pull/499) Bugfix: MPI datatype mapping: `torch.int16` now maps to `MPI.SHORT` instead of `MPI.SHORT_INT`
+- [#501](https://github.com/helmholtz-analytics/heat/pull/501) New Feature: flatten
 - [#506](https://github.com/helmholtz-analytics/heat/pull/506) Bugfix: setup.py has correct version parsing
 - [#507](https://github.com/helmholtz-analytics/heat/pull/507) Bugfix: sanitize_axis changes axis of 0-dim scalars to None
 - [#515](https://github.com/helmholtz-analytics/heat/pull/515) ht.var() now returns the unadjusted sample variance by default, Bessel's correction can be applied by setting ddof=1.

heat/core/dndarray.py

@@ -1060,6 +1060,23 @@ def expand_dims(self, axis):
         """
         return manipulations.expand_dims(self, axis)
 
+    def flatten(self):
+        """
+        Return a flat tensor.
+
+        Returns
+        -------
+        flattened : ht.DNDarray
+            The flattened tensor
+
+        Examples
+        --------
+        >>> x = ht.array([[1,2],[3,4]])
+        >>> x.flatten()
+        tensor([1,2,3,4])
+        """
+        return manipulations.flatten(self)
+
     def __float__(self):
         """
         Float scalar casting.

heat/core/manipulations.py

@@ -1,5 +1,6 @@
 import numpy as np
 import torch
+import warnings
 
 from .communication import MPI
 
@@ -15,6 +16,7 @@
     "diag",
     "diagonal",
     "expand_dims",
+    "flatten",
     "flip",
     "flipud",
     "hstack",
@@ -594,6 +596,49 @@ def expand_dims(a, axis):
     )
 
 
+def flatten(a):
+    """
+    Flattens an array into one dimension.
+    WARNING: if a.split > 0, then the array must be resplit.
+
+    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,
+        )
+
+    if a.split > 0:
+        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 flip(a, axis=None):
     """
     Reverse the order of elements in an array along the given axis.

heat/core/tests/test_dndarray.py

@@ -164,6 +164,14 @@ def test_complex_cast(self):
             with self.assertRaises(TypeError):
                 complex(ht.full((ht.MPI_WORLD.size,), 2, split=0, device=ht_device))
 
+    def test_flatten(self):
+        a = ht.ones((4, 4, 4), split=1)
+        result = ht.ones((64,), split=0)
+        flat = a.flatten()
+
+        self.assertEqual(flat.shape, result.shape)
+        self.assertTrue(ht.equal(flat, result))
+
     def test_fill_diagonal(self):
         ref = ht.zeros(
             (ht.MPI_WORLD.size * 2, ht.MPI_WORLD.size * 2),