review: apply @mloubout's suggestions

* use ``devito.tools.as_tuple`` to normalize the polymorphic ``axes`` of
  ``transpose`` (no-args, single ``None``, single tuple/list, or
  individual positional args all collapse to a flat tuple).
* drop the ``isinstance(ret, Data)`` guards in ``transpose`` and
  ``swapaxes`` -- ``numpy.ndarray.<view-op>`` preserves the subclass via
  ``__array_finalize__``, so ``ret`` is always a ``Data``.
* replace the hand-built reference in ``test_slice_after_transpose``
  with ``ref = np.array(f.data)`` (and the 3D equivalent for ``g``),
  removing the duplicated ``arange/reshape/astype`` lines; shape and
  values are still compared against the NumPy reference array via
  ``np.array_equal``.

Author: gaoflow

devito/data/data.py

@@ -220,20 +220,19 @@ def transpose(self, *axes):
         wrong per-axis decomposition and silently returns a wrong-shaped
         result (see issue #2187).
         """
-        if not axes or axes == (None,):
-            new_order = tuple(range(self.ndim - 1, -1, -1))
-        elif len(axes) == 1 and isinstance(axes[0], Iterable):
-            new_order = tuple(axes[0])
-        else:
-            new_order = tuple(axes)
-        # Normalize negative axes (numpy accepts them)
-        new_order = tuple(ax % self.ndim for ax in new_order)
+        # Accept the same axis-spec forms as ``numpy.ndarray.transpose``:
+        # no args, a single ``None``, a single tuple/list, or per-arg.
+        if len(axes) == 1:
+            axes = as_tuple(axes[0])
+        new_order = (
+            tuple(range(self.ndim - 1, -1, -1)) if not axes
+            else tuple(ax % self.ndim for ax in axes)
+        )
 
         ret = super().transpose(*axes)
-        if isinstance(ret, Data):
-            ret._decomposition = tuple(self._decomposition[i] for i in new_order)
-            ret._modulo = tuple(self._modulo[i] for i in new_order)
-            ret._is_distributed = any(d is not None for d in ret._decomposition)
+        ret._decomposition = tuple(self._decomposition[i] for i in new_order)
+        ret._modulo = tuple(self._modulo[i] for i in new_order)
+        ret._is_distributed = any(d is not None for d in ret._decomposition)
         return ret
 
     def swapaxes(self, axis1, axis2):
@@ -245,12 +244,11 @@ def swapaxes(self, axis1, axis2):
         axis1 = axis1 % self.ndim
         axis2 = axis2 % self.ndim
         ret = super().swapaxes(axis1, axis2)
-        if isinstance(ret, Data):
-            order = list(range(self.ndim))
-            order[axis1], order[axis2] = order[axis2], order[axis1]
-            ret._decomposition = tuple(self._decomposition[i] for i in order)
-            ret._modulo = tuple(self._modulo[i] for i in order)
-            ret._is_distributed = any(d is not None for d in ret._decomposition)
+        order = list(range(self.ndim))
+        order[axis1], order[axis2] = order[axis2], order[axis1]
+        ret._decomposition = tuple(self._decomposition[i] for i in order)
+        ret._modulo = tuple(self._modulo[i] for i in order)
+        ret._is_distributed = any(d is not None for d in ret._decomposition)
         return ret
 
     @property

tests/test_data.py

@@ -223,44 +223,31 @@ def test_slice_after_transpose(self):
         grid = Grid(shape=(4, 6))
         f = Function(name='f', grid=grid)
         f.data[:] = np.arange(24).reshape((4, 6)).astype(np.float32)
-
-        ref = np.arange(24).reshape((4, 6)).astype(np.float32)
+        ref = np.array(f.data)
 
         # ``.T`` (C-level shortcut) then slice
-        out = f.data.T[::2, ::2]
-        assert out.shape == ref.T[::2, ::2].shape
-        assert np.array_equal(np.asarray(out), ref.T[::2, ::2])
+        assert np.array_equal(f.data.T[::2, ::2], ref.T[::2, ::2])
 
         # Equivalent: slice then ``.T``
-        out2 = f.data[::2, ::2].T
-        assert out2.shape == ref[::2, ::2].T.shape
-        assert np.array_equal(np.asarray(out2), ref[::2, ::2].T)
+        assert np.array_equal(f.data[::2, ::2].T, ref[::2, ::2].T)
 
         # Explicit ``transpose`` call -- same behavior as ``.T``
-        out3 = f.data.transpose()[::2, ::2]
-        assert out3.shape == ref.transpose()[::2, ::2].shape
-        assert np.array_equal(np.asarray(out3), ref.transpose()[::2, ::2])
+        assert np.array_equal(f.data.transpose()[::2, ::2],
+                              ref.transpose()[::2, ::2])
 
         # ``swapaxes`` between non-conforming dims
-        out4 = f.data.swapaxes(0, 1)[::2, ::2]
-        assert out4.shape == ref.swapaxes(0, 1)[::2, ::2].shape
-        assert np.array_equal(np.asarray(out4), ref.swapaxes(0, 1)[::2, ::2])
+        assert np.array_equal(f.data.swapaxes(0, 1)[::2, ::2],
+                              ref.swapaxes(0, 1)[::2, ::2])
 
         # 3D transpose with an explicit axis order, then per-axis slice
         grid3 = Grid(shape=(2, 4, 6))
         g = Function(name='g3', grid=grid3)
         g.data[:] = np.arange(48).reshape((2, 4, 6)).astype(np.float32)
-        ref3 = np.arange(48).reshape((2, 4, 6)).astype(np.float32)
-
-        out5 = g.data.T[::2, ::2, ::2]
-        assert out5.shape == ref3.T[::2, ::2, ::2].shape
-        assert np.array_equal(np.asarray(out5), ref3.T[::2, ::2, ::2])
+        ref3 = np.array(g.data)
 
-        out6 = g.data.transpose((1, 0, 2))[::2, ::1, ::3]
-        assert out6.shape == ref3.transpose((1, 0, 2))[::2, ::1, ::3].shape
-        assert np.array_equal(
-            np.asarray(out6), ref3.transpose((1, 0, 2))[::2, ::1, ::3]
-        )
+        assert np.array_equal(g.data.T[::2, ::2, ::2], ref3.T[::2, ::2, ::2])
+        assert np.array_equal(g.data.transpose((1, 0, 2))[::2, ::1, ::3],
+                              ref3.transpose((1, 0, 2))[::2, ::1, ::3])
 
     def test_transpose_permutes_data_metadata(self):
         """