features/239-flip()

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"
+- [#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`
 
 # v0.3.0

heat/core/manipulations.py

@@ -14,6 +14,7 @@
     "diag",
     "diagonal",
     "expand_dims",
+    "flip",
     "hstack",
     "resplit",
     "sort",
@@ -556,6 +557,68 @@ def expand_dims(a, axis):
     )
 
 
+def flip(a, axis=None):
+    """
+    Reverse the order of elements in an array along the given axis.
+
+    The shape of the array is preserved, but the elements are reordered.
+
+    Parameters
+    ----------
+    a: ht.DNDarray
+        Input array to be flipped
+    axis: int, tuple
+        A list of axes to be flipped
+
+    Returns
+    -------
+    res: ht.DNDarray
+        The flipped array.
+
+    Examples
+    --------
+    >>> a = ht.array([[0,1],[2,3]])
+    >>> ht.flip(a, [0])
+    tensor([[2, 3],
+        [0, 1]])
+
+    >>> b = ht.array([[0,1,2],[3,4,5]], split=1)
+    >>> ht.flip(a, [0,1])
+    (1/2) tensor([5,4,3])
+    (2/2) tensor([2,1,0])
+    """
+    # flip all dimensions
+    if axis is None:
+        axis = tuple(range(a.numdims))
+
+    # torch.flip only accepts tuples
+    if isinstance(axis, int):
+        axis = [axis]
+
+    flipped = torch.flip(a._DNDarray__array, axis)
+
+    if a.split not in axis:
+        return factories.array(
+            flipped, dtype=a.dtype, is_split=a.split, device=a.device, comm=a.comm
+        )
+
+    # Need to redistribute tensors on split axis
+    # Get local shapes
+    old_lshape = a.lshape
+    dest_proc = a.comm.size - 1 - a.comm.rank
+    new_lshape = a.comm.sendrecv(old_lshape, dest=dest_proc, source=dest_proc)
+
+    # Exchange local tensors
+    req = a.comm.Isend(flipped, dest=dest_proc)
+    received = torch.empty(new_lshape, dtype=a._DNDarray__array.dtype, device=a.device.torch_device)
+    a.comm.Recv(received, source=dest_proc)
+
+    res = factories.array(received, dtype=a.dtype, is_split=a.split, device=a.device, comm=a.comm)
+    res.balance_()  # after swapping, first processes may be empty
+    req.Wait()
+    return res
+
+
 def hstack(tup):
     """
     Stack arrays in sequence horizontally (column wise).

heat/core/tests/test_manipulations.py

@@ -763,6 +763,29 @@ def test_expand_dims(self):
         with self.assertRaises(ValueError):
             ht.empty((3, 4, 5), device=ht_device).expand_dims(-5)
 
+    def test_flip(self):
+        a = ht.array([1, 2], device=ht_device)
+        r_a = ht.array([2, 1], device=ht_device)
+        self.assertTrue(ht.equal(ht.flip(a, 0), r_a))
+
+        a = ht.array([[1, 2], [3, 4]], device=ht_device)
+        r_a = ht.array([[4, 3], [2, 1]], device=ht_device)
+        self.assertTrue(ht.equal(ht.flip(a), r_a))
+
+        a = ht.array([[2, 3], [4, 5], [6, 7], [8, 9]], split=1, dtype=ht.float32, device=ht_device)
+        r_a = ht.array(
+            [[9, 8], [7, 6], [5, 4], [3, 2]], split=1, dtype=ht.float32, device=ht_device
+        )
+        self.assertTrue(ht.equal(ht.flip(a, [0, 1]), r_a))
+
+        a = ht.array(
+            [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], split=0, dtype=ht.uint8, device=ht_device
+        )
+        r_a = ht.array(
+            [[[3, 2], [1, 0]], [[7, 6], [5, 4]]], split=0, dtype=ht.uint8, device=ht_device
+        )
+        self.assertTrue(ht.equal(ht.flip(a, [1, 2]), r_a))
+
     def test_hstack(self):
         # cases to test:
         # MM===================================