add test cases

Author: FynnBe

tests/test_tensor.py

@@ -1,3 +1,5 @@
+from typing import Sequence
+
 import numpy as np
 import pytest
 import xarray as xr
@@ -8,9 +10,19 @@
 
 @pytest.mark.parametrize(
     "axes",
-    ["yx", "xy", "cyx", "yxc", "bczyx", "xyz", "xyzc", "bzyxc"],
+    [
+        "yx",
+        "xy",
+        "cyx",
+        "yxc",
+        "bczyx",
+        "xyz",
+        "xyzc",
+        "bzyxc",
+        ("batch", "channel", "x", "y"),
+    ],
 )
-def test_transpose_tensor_2d(axes: str):
+def test_transpose_tensor_2d(axes: Sequence[str]):
 
     tensor = Tensor.from_numpy(np.random.rand(256, 256), dims=None)
     transposed = tensor.transpose([AxisId(a) for a in axes])
@@ -19,9 +31,18 @@ def test_transpose_tensor_2d(axes: str):
 
 @pytest.mark.parametrize(
     "axes",
-    ["zyx", "cyzx", "yzixc", "bczyx", "xyz", "xyzc", "bzyxtc"],
+    [
+        "zyx",
+        "cyzx",
+        "yzixc",
+        "bczyx",
+        "xyz",
+        "xyzc",
+        "bzyxtc",
+        ("batch", "channel", "x", "y", "z"),
+    ],
 )
-def test_transpose_tensor_3d(axes: str):
+def test_transpose_tensor_3d(axes: Sequence[str]):
     tensor = Tensor.from_numpy(np.random.rand(64, 64, 64), dims=None)
     transposed = tensor.transpose([AxisId(a) for a in axes])
     assert transposed.ndim == len(axes)