Implement moveaxis

Author: Ricardo Vieira

aesara/tensor/basic.py

@@ -10,11 +10,14 @@
 from collections.abc import Sequence
 from functools import partial
 from numbers import Number
-from typing import Optional, Tuple, Union
+from typing import Optional
+from typing import Sequence as TypeSequence
+from typing import Tuple, Union
 from typing import cast as type_cast
 
 import numpy as np
 from numpy.core.multiarray import normalize_axis_index
+from numpy.core.numeric import normalize_axis_tuple
 
 import aesara
 import aesara.scalar.sharedvar
@@ -3635,6 +3638,51 @@ def swapaxes(y, axis1, axis2):
     return y.dimshuffle(li)
 
 
+def moveaxis(
+    a: Union[np.ndarray, TensorVariable],
+    source: Union[int, TypeSequence[int]],
+    destination: Union[int, TypeSequence[int]],
+) -> TensorVariable:
+    """Move axes of a TensorVariable to new positions.
+
+    Other axes remain in their original order.
+
+    Parameters
+    ----------
+    a
+        The TensorVariable whose axes should be reordered.
+    source
+        Original positions of the axes to move. These must be unique.
+    destination
+        Destination positions for each of the original axes. These must also be
+        unique.
+
+    Returns
+    -------
+    result
+        TensorVariable with moved axes.
+
+    """
+
+    a = as_tensor_variable(a)
+
+    source = normalize_axis_tuple(source, a.ndim, "source")
+    destination = normalize_axis_tuple(destination, a.ndim, "destination")
+
+    if len(source) != len(destination):
+        raise ValueError(
+            "`source` and `destination` arguments must have the same number of elements"
+        )
+
+    order = [n for n in range(a.ndim) if n not in source]
+
+    for dest, src in sorted(zip(destination, source)):
+        order.insert(dest, src)
+
+    result = a.dimshuffle(order)
+    return result
+
+
 def choose(a, choices, mode="raise"):
     """
     Construct an array from an index array and a set of arrays to choose from.
@@ -4014,6 +4062,7 @@ def take_along_axis(arr, indices, axis=0):
     "atleast_3d",
     "choose",
     "swapaxes",
+    "moveaxis",
     "stacklists",
     "diag",
     "diagonal",

tests/tensor/test_basic.py

@@ -60,6 +60,7 @@
     join,
     make_vector,
     mgrid,
+    moveaxis,
     nonzero,
     nonzero_values,
     ogrid,
@@ -3984,6 +3985,23 @@ def test_numpy_compare(self):
         assert np.allclose(n_s, t_s)
 
 
+def test_moveaxis():
+    x = at.zeros((3, 4, 5))
+    tuple(moveaxis(x, 0, -1).shape.eval()) == (4, 5, 3)
+    tuple(moveaxis(x, -1, 0).shape.eval()) == (5, 3, 4)
+    tuple(moveaxis(x, [0, 1], [-1, -2]).shape.eval()) == (5, 4, 3)
+    tuple(moveaxis(x, [0, 1, 2], [-1, -2, -3]).shape.eval()) == (5, 4, 3)
+
+
+def test_moveaxis_error():
+    x = at.zeros((3, 4, 5))
+    with pytest.raises(
+        ValueError,
+        match="`source` and `destination` arguments must have the same number of elements",
+    ):
+        moveaxis(x, [0, 1], 0)
+
+
 class TestChoose(utt.InferShapeTester):
     op = staticmethod(choose)
     op_class = Choose