Skip to content

fix: permute Data._decomposition and _modulo on transpose (fixes #2187)#2938

Open
gaoflow wants to merge 2 commits into
devitocodes:mainfrom
gaoflow:fix/data-transpose-decomposition-modulo
Open

fix: permute Data._decomposition and _modulo on transpose (fixes #2187)#2938
gaoflow wants to merge 2 commits into
devitocodes:mainfrom
gaoflow:fix/data-transpose-decomposition-modulo

Conversation

@gaoflow
Copy link
Copy Markdown

@gaoflow gaoflow commented May 28, 2026

Fix: permute Data._decomposition / _modulo on transpose

Fixes #2187.

Problem

A Data view created via .T, transpose(...) or swapaxes(...)
inherited the source's _decomposition and _modulo tuples unchanged
through __array_finalize__ (the self.ndim == obj.ndim branch
explicitly does self._modulo = obj._modulo,
self._decomposition = obj._decomposition). Both tuples are indexed by
axis, so after the axes are permuted they point at the wrong
dimensions. A subsequent slice goes through __getitem__
_index_glb_to_loc(glb_idx), which uses the (stale) _decomposition
to translate the global slice into local-axis ranges, and silently
returns a wrong-shaped view.

The minimal example from #2187:

import numpy as np
from devito import Grid, Function

grid = Grid(shape=(4, 6))
f = Function(name='f', grid=grid)

# NumPy reference -- expected shape (3, 2):
np.zeros((4, 6)).T[::2, ::2].shape   # (3, 2)

# Devito on `main` -- wrong shape:
f.data.T[::2, ::2].shape             # (2, 2)

f.data[::2, ::2].T (slice then transpose) works because the slice
goes through __getitem__ first and the second-step transpose only
flips a 2-element metadata tuple whose entries are already aligned.

Why both transpose() and T need overriding

ndarray.T is a C-level shortcut: it does not dispatch through
Python's transpose method on a subclass. Verified directly:

class A(np.ndarray):
    def transpose(self, *axes): print("called!"); return super().transpose(*axes)

a = np.zeros((3, 4)).view(A)
a.T              # silent -- C path
a.transpose()    # prints "called!"
np.transpose(a)  # prints "called!"

So overriding transpose alone leaves .T buggy; overriding T alone
leaves explicit transpose / swapaxes calls buggy. Three thin
overrides are added on Data:

  • transpose(*axes) — normalize axes (no-arg, single
    tuple/list/iterable, or per-arg individual axes), call
    super().transpose, then permute _decomposition / _modulo into
    the new axis order and refresh _is_distributed so it stays
    consistent.
  • swapaxes(axis1, axis2) — same treatment for the two-axis case.
  • T — explicit Python property that forwards to transpose().

Verification

The two new tests in tests/test_data.py::TestDataBasic:

  • test_slice_after_transpose — covers .T, transpose(),
    transpose((1, 0, 2)), swapaxes(0, 1), in both 2D and 3D, with
    the NumPy reference array as the source of truth (shape and
    values are compared).
  • test_transpose_permutes_data_metadata — asserts
    _decomposition / _modulo end up in the new order (reverse for
    .T, swap for swapaxes, explicit permutation for
    transpose((1, 2, 0))).

Local run on main HEAD (no MPI installed; MPI-only tests pre-existing
failures are unchanged):

$ pytest tests/test_data.py -v -k "not gather and not singlempi and not Distributed and not mpi"
============ 37 passed, 2 skipped, ...
$ pytest tests/test_pickle.py tests/test_dimension.py -k "not parallel"
============ 600 passed, ...
$ ruff check devito/data/data.py tests/test_data.py
All checks passed!

Scope

Just devito/data/data.py (production fix) and tests/test_data.py
(regression tests). No public API change — transpose / T /
swapaxes signatures are unchanged; behavior matches NumPy.

A ``Data`` view created by ``.T`` (or ``transpose``/``swapaxes``)
inherited the source's ``_decomposition`` and ``_modulo`` tuples
unchanged via ``__array_finalize__``. Both tuples are indexed by
*axis*, so after the axes are permuted they point at the wrong
dimensions. A subsequent ``__getitem__`` then routes the slice
through ``_index_glb_to_loc`` against the wrong per-axis ranges
and silently returns a wrong-shaped result, e.g.:

    grid = Grid(shape=(4, 6))
    f = Function(name='f', grid=grid)
    f.data.T[::2, ::2].shape  # (2, 2) -- wrong, NumPy gives (3, 2)

Overriding ``.T`` alone is not enough because ``ndarray.T`` is a
C-level shortcut that does not dispatch through Python's
``transpose``; ``transpose`` alone is not enough either because
``.T`` then keeps the old (buggy) behavior. So three thin
overrides are added on ``Data``:

* ``transpose(*axes)``: normalize ``axes`` (no-arg, sequence, or
  per-arg), call ``super().transpose``, then permute
  ``_decomposition``/``_modulo`` into the new axis order and
  refresh ``_is_distributed`` so it stays consistent.
* ``swapaxes(axis1, axis2)``: same treatment for the two-axis case.
* ``T``: explicit Python property that forwards to ``transpose``.

Slice-then-transpose (``f.data[::2, ::2].T``) already worked
because the slice goes through ``__getitem__`` first and the
transpose only flips a 2-element metadata tuple where both entries
match. The fix preserves that path unchanged.

Adds ``test_slice_after_transpose`` (covers ``.T``, ``transpose``,
``swapaxes``, 2D and 3D, with the NumPy reference array as the
source of truth) and ``test_transpose_permutes_data_metadata``
(asserts ``_decomposition``/``_modulo`` end up in the new order).

Fixes devitocodes#2187.
Copy link
Copy Markdown
Contributor

@mloubout mloubout left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for that, looks sane and definitely a bug. Some small comments

Comment thread devito/data/data.py Outdated
"""
if not axes or axes == (None,):
new_order = tuple(range(self.ndim - 1, -1, -1))
elif len(axes) == 1 and isinstance(axes[0], Iterable):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can use devito.tools.as_tuple for these two cases

Comment thread devito/data/data.py Outdated
new_order = tuple(ax % self.ndim for ax in new_order)

ret = super().transpose(*axes)
if isinstance(ret, Data):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this ever False?

Comment thread devito/data/data.py Outdated
axis1 = axis1 % self.ndim
axis2 = axis2 % self.ndim
ret = super().swapaxes(axis1, axis2)
if isinstance(ret, Data):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same

Comment thread tests/test_data.py Outdated
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)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ref = np.array(f.data)

* 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``.
@mloubout
Copy link
Copy Markdown
Contributor

Thank you, please rebase with commit messages according to guidlines and this is gtg

https://github.com/devitocodes/devito?tab=contributing-ov-file#commit-messages-and-pull-request-titles

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Slicing Data returns incorrect results following a transpose()

2 participants