ENH: Introduce moveaxis function (#99)

* ENH: Introduce `moveaxis` function

* Bump version

Author: mtsokol

pixi.toml

@@ -4,7 +4,7 @@ channels = ["conda-forge"]
 description = "Add a short description here"
 name = "finch-tensor"
 platforms = ["osx-arm64"]
-version = "0.2.9"
+version = "0.2.10"
 
 [tasks]
 compile = "python -c 'import finch'"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "finch-tensor"
-version = "0.2.9"
+version = "0.2.10"
 description = ""
 authors = ["Willow Ahrens <[email protected]>"]
 readme = "README.md"

src/finch/__init__.py

@@ -52,6 +52,7 @@
     eye,
     tensordot,
     permute_dims,
+    moveaxis,
     where,
     nonzero,
     sum,
@@ -166,6 +167,7 @@
     "tensordot",
     "matmul",
     "permute_dims",
+    "moveaxis",
     "where",
     "nonzero",
     "int_",

src/finch/tensor.py

@@ -868,11 +868,19 @@ def linspace(
     )
 
 
-def permute_dims(x: Tensor, axes: tuple[int, ...]):
+def permute_dims(x: Tensor, axes: tuple[int, ...]) -> Tensor:
     return x.permute_dims(axes)
 
 
-def astype(x: Tensor, dtype: DType, /, *, copy: bool = True):
+def moveaxis(x: Tensor, source: int, destination: int) -> Tensor:
+    axes = list(range(x.ndim))
+    norm_source = normalize_axis_index(source, x.ndim)
+    norm_dest = normalize_axis_index(destination, x.ndim)
+    axes.insert(norm_dest, axes.pop(norm_source))
+    return x.permute_dims(tuple(axes))
+
+
+def astype(x: Tensor, dtype: DType, /, *, copy: bool = True) -> Tensor:
     if not copy:
         if x.dtype == dtype:
             return x

tests/test_sparse.py

@@ -164,6 +164,20 @@ def test_permute_dims(arr3d, permutation, format, order, opt):
     assert_equal(actual_lazy_mode.todense(), expected)
 
 
+@pytest.mark.parametrize(
+    "src_dest", [(0, 1), (1, 0), (-1, 2), (-2, -1), (1, 1)]
+)
+@parametrize_optimizer
+def test_moveaxis(arr3d, src_dest, opt):
+    finch.set_optimizer(opt)
+    src, dest = src_dest
+    arr_finch = finch.Tensor(arr3d)
+
+    actual = finch.moveaxis(arr_finch, src, dest)
+    expected = np.moveaxis(arr3d, src, dest)
+    assert_equal(actual.todense(), expected)
+
+
 @pytest.mark.parametrize("order", ["C", "F"])
 @parametrize_optimizer
 def test_astype(arr3d, order, opt):