Skip to content

Feat: WIP - Enable ND-Distributed arrays on operators #198

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mrava87 wants to merge 10 commits into
base: main
Choose a base branch
from
Draft

Feat: WIP - Enable ND-Distributed arrays on operators #198

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
4b131dc
WIP: enable ND-Distributed arrays on operators
mrava87SW Jun 2, 2026
f80103d
Update reshaped decorator and dot method
rohanbabbar04 Jun 23, 2026
2286152
Update method to handle different axis
rohanbabbar04 Jun 23, 2026
6285e30
Update reshaped and matvec sections
rohanbabbar04 Jul 3, 2026
5076b05
Add local_array, and update ravel, reshape and reshaped decorator
rohanbabbar04 Jul 5, 2026
9040495
Add _local_dims_dimsd_axis and Update First and second derivative
rohanbabbar04 Jul 6, 2026
ce59583
Add local_dims and local_dimsd, update First,SecondDerivate, and Bloc…
rohanbabbar04 Jul 7, 2026
95f8fc7
Update __add__ dunder
rohanbabbar04 Jul 8, 2026
8452b9c
Remove self._local_dims_dimsd
rohanbabbar04 Jul 8, 2026
06b3632
Update docstrings of reshape method
rohanbabbar04 Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 47 additions & 6 deletions pylops_mpi/DistributedArray.py
View file
Open in desktop

@mrava87 mrava87 Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Nice idea about the local_array

Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,10 @@ class DistributedArray(DistributedMixIn):
Engine used to store array (``numpy`` or ``cupy``)
dtype : :obj:`str`, optional
Type of elements in input array. Defaults to ``numpy.float64``.
local_array : :obj:`numpy.ndarray` or :obj:`cupy.ndarray`, optional

@mrava87 mrava87 Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Since most of the inputs are positional, I would move this a bit earlier maybe close to local_shapes... and a few things to consider when one passes local_array:

  • dtype can be conflicting, so maybe we should say if local_array is passed dtype is ignored and local_array.dtype used instead?
  • same for engine?
  • similarly local_shapes may be contrasting, maybe here I would raise an error if both are passed to avoid confusion?

I understand users may not find this useful (aka it is more for developers) but since it surface as an input of the class constructor, I think we should treat it as is and make sure it does not lead to confusing scenarios...

Existing local array to use as the underlying storage instead of
allocating a new one. The shape of the local_array must match the local
shape at that rank.
"""

def __init__(self, global_shape: Union[Tuple, Integral],
Expand All @@ -149,7 +153,8 @@ def __init__(self, global_shape: Union[Tuple, Integral],
local_shapes: Optional[List[Union[Tuple, Integral]]] = None,
mask: Optional[List[Integral]] = None,
engine: Optional[str] = "numpy",
dtype: Optional[DTypeLike] = np.float64):
dtype: Optional[DTypeLike] = np.float64,
local_array: Optional[NDArray] = None):
if isinstance(global_shape, Integral):
global_shape = (global_shape,)
if len(global_shape) <= axis:
Expand Down Expand Up @@ -179,7 +184,18 @@ def __init__(self, global_shape: Union[Tuple, Integral],
partition, axis)

self._engine = engine
self._local_array = get_module(engine).empty(shape=self.local_shape, dtype=self.dtype)
if local_array is None:
self._local_array = get_module(engine).empty(
shape=self.local_shape,
dtype=self.dtype,
)
else:
if local_array.shape != self.local_shape:
raise ValueError(
f"local_array has shape {local_array.shape}, "
f"expected {self.local_shape}"
)
self._local_array = local_array

def __getitem__(self, index):
return self.local_array[index]
Expand Down Expand Up @@ -851,17 +867,42 @@ def ravel(self, order: Optional[str] = "C"):
Flattened 1-D DistributedArray
"""
local_shapes = [(np.prod(local_shape, axis=-1), ) for local_shape in self.local_shapes]
local_array = np.ravel(self.local_array, order=order)
arr = DistributedArray(global_shape=np.prod(self.global_shape),
base_comm=self.base_comm,
base_comm_nccl=self.base_comm_nccl,
local_shapes=local_shapes,
mask=self.mask,
partition=self.partition,
engine=self.engine,
dtype=self.dtype)
local_array = np.ravel(self.local_array, order=order)
x = local_array.copy()
arr[:] = x
dtype=self.dtype,
local_array=local_array)
return arr

def reshape(self, local_shape: Tuple, axis: Optional[int] = 0):
"""Return a reshaped DistributedArray

Parameters
----------
local_shape : :obj:`tuple`
Shape of the local array on each MPI rank.
axis : :obj:`int`, optional
Distribution axis in the reshaped global array. Defaults to ``0``.
"""
local_shapes = self.base_comm.allgather(local_shape)
global_shape = list(local_shapes[0])

@mrava87 mrava87 Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Does this assume that if the array is scattered all dims but the one of axis are the same? Do you think its worth a check+raise or we shall just assume the user does it correctly (at least maybe mention this in the docstring?)

if self.partition is Partition.SCATTER:
global_shape[axis] = np.sum([ls[axis] for ls in local_shapes])
local_array = self.local_array.reshape(local_shapes[self.rank])
arr = DistributedArray(global_shape=tuple(global_shape),
base_comm=self.base_comm,
base_comm_nccl=self.base_comm_nccl,
local_shapes=local_shapes,
mask=self.mask,
partition=self.partition,
engine=self.engine,
dtype=self.dtype,
local_array=local_array)
return arr

def empty_like(self):
Expand Down
50 changes: 41 additions & 9 deletions pylops_mpi/LinearOperator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from scipy.sparse._sputils import isintlike, isshape
from scipy.sparse.linalg._interface import _get_dtype

from pylops import LinearOperator
from pylops import LinearOperator, get_ndarray_multiplication
from pylops.utils import DTypeLike, ShapeLike

from pylops_mpi import DistributedArray
Expand Down Expand Up @@ -259,9 +259,11 @@ def dot(self, x):
Op = _ProductLinearOperator(self, x)
self._copy_attributes(
Op,
exclude=['dims']
exclude=['dims', '_local_dims']

@mrava87 mrava87 Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

So this _local_dims is not set in MPILinearOperator constructor, so if one does MPILinearOperator(Op) with Op being a PyLops operator, this will not be present... thinking if this could lead to any scenario where it is always expected but it may not always be there?

@mrava87 mrava87 Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

For example here you don't set _local_dimsd ?

)
Op.dims = x.dims
if getattr(x, "_local_dims", None) is not None:
Op._local_dims = x._local_dims
return Op
elif np.isscalar(x):
Op = _ScaledLinearOperator(self, x)
Expand All @@ -270,11 +272,29 @@ def dot(self, x):
)
return Op
else:
if x is None or x.ndim == 1:
return self.matvec(x)
if not get_ndarray_multiplication() and x.ndim >= 2:
msg = (
"Operator can only be applied to 1D vectors. "
"Enable ndarray multiplication with pylops.set_ndarray_multiplication(True)."
)
raise ValueError(msg)
is_dims_shaped = x.global_shape == self.dims
if is_dims_shaped:
x = x.redistribute(axis=0).ravel()

@mrava87 mrava87 Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Mmh back to one of my original comments: why do we need ravel here, and is with the new local_array now a much cheaper (view-only) operation?

@mrava87 mrava87 Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

But since redistributed is for sure not cheap, why you see this as needed here? And why ravel, other than maybe to be as close as possible to what we do in pylops?

@rohanbabbar04 rohanbabbar04 Jul 11, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

@mrava87 My main thought was to try and mimic what we do in pylops, which was flattening and then apply the _matvec and then reshape back to dimsd like we do in pylops. With the introduction of local_array, copy would be cheap, but yeah with redistribution it will be expensive. So would you suggest we keep NDArray only for specific ops like FFTs(and maybe other....), since we used to redistribute(expensive operation) it to axis 0 to apply the ravel at the end.

@mrava87 mrava87 Jul 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

@rohanbabbar04 I am not sure, my initial thought was 'let's follow the PyLops steps' but then already when I started toying around I saw the issue with us recreating and reassigning all of these numpy arrays (as our local arrays) so I had the idea that maybe for the purpose of pylops-mpi (solve one very large problem in a distributed manner*) we could just avoid doing any ravel-(reshape-apply-flattened)-reshape with (reshape-apply-flattened) done in the reshape decorator if an input was not a 1d array.

Now if this process was cheap as it seems now (up to redistribute) this would be purely stylistical but given the fact we do that, maybe we need to rethink the entire process instead of fitting the PyLops pattern in it?

  • for pylops the idea of solving one or more flattened problems (with multiple RHS) makes somehow sense as in some applications one has a small problem, but many of them... this kind of goes away in pylops-mpi as if you reach for it the likely cause is your problem is so big it cant even fit in one single machine VRAM (or is very slow when solved in single-cpu mode), so my thought that maybe we can loosely allow a solver to operate with either flat or unflattened arrays just considering them all the time as if they were flattened (where basically only for dot and norm I see we needed some modifications, all +/-/* are already transparent to that)
. So would you suggest we keep NDArray only for specific ops like FFTs(and maybe other....), since we used to redistribute(expensive operation) it to axis 0 to apply the ravel at the end.

but what I pointed here is applied to all operators where one passes a ND-array and self.dims is not 1d, so I am not sure what you say would work? Perhaps I would suggest we take a tutorial, say poststack and see if what we have now (plus likely some small changes to the solvers) would allow going from doing

# Create flattened model data
m3d_dist = pylops_mpi.DistributedArray(global_shape=ny * nx * nz)
m3d_dist[:] = m3d_i.flatten()

to

# Create distributed model data
m3d_dist = pylops_mpi.DistributedArray(global_shape=(ny, nx, nz), axis=0)
m3d_dist[:] = m3d_i

where note that here this is kind of just stylistic and for ease of use (I always thought it would be nice for a user not to think in 1d distributed flattened arrays, but for other scenarios where an operator like FFT is part of the final Op it would become more important.

Right now I have the feeling that even doing d_dist = BDiag @ m3d_dist to the unflatted m3d_dist would hit the heavy redistribute?

@rohanbabbar04 rohanbabbar04 Jul 14, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Hmm... Yeah, I was also thinking about the best way to tackle ND arrays, especially for operators like FFTs and derivatives where users would naturally want to preserve the multidimensional shape.

I was thinking of changing matvec slightly:

def matvec(self, x: DistributedArray) -> DistributedArray:
    M, N = self.shape
    if x.ndim == 1:
        if x.global_shape != (N,):
            raise ValueError("dimension mismatch")
        return self._matvec(x)

    if np.prod(x.global_shape) != N:
        raise ValueError("dimension mismatch")

    return self._matvec_nd(x)

The idea would be that _matvec_nd is optional. Operators that make more sense/reduce communication work with ND arrays (such as FFTs or derivatives) can implement it and preserve the input shape. Operators that don't have a meaningful ND implementation can simply leave _matvec_nd unimplemented (e.g., raising NotImplementedError) and continue to support only 1D distributed arrays through _matvec.

For FFTs, for example, we could move the implementation entirely into _matvec_nd, since the FFT is expensive when we are creating a 1-d array and the users would typically want to work with the ND dimensions rather than flattening the array.

Again, this is just an idea for now, and we can think it through some more. 🙂

@mrava87 mrava87 Jul 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Right, the idea seems appealing but l would need to know a bit more to see if this makes sense 😉

First question: as we want to mantain backward compatibility, I guess for operators where it makes sense to implement _matvec_nd their _matvec would still leverage the @reshape decorator and then just call _matvec_nd?

Second question: this is something I am always very careful we honor, as I think it's one of the great features of pylops. Users should always be able to mix any combination of operators (provided shapes make sense) in the most transparent possible way (i.e., not having to know all the internals of the actual implementations of the operators). Now let' say we go for this _matvec_nd implementation, do you think we can end up in a situation where:
- a user can always create a combined operator and pass a 1D distributed array, even if somehow performance-wise suboptimal, like when you have an FFT operatr
- a user can always create a combined operator and pass a ND distributed array provided all of the operators in the combined operator support this (implement _matvec_nd)?

Again, I would maybe start by taking a tutorial like I suggested above and see if any of what you are thinking would enable/not-enable solving that problem with ND distributed arrays, and then maybe work backwards adding support to more operators and checking against other tutorials/examples

@rohanbabbar04 rohanbabbar04 Jul 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Yup, sounds good, we can start with a tutorial and see the best way possible 🙂.

if x.ndim == 1:
y = self.matvec(x)
if is_dims_shaped and get_ndarray_multiplication():
_local_dimsd = getattr(self, "_local_dimsd", None)
if _local_dimsd:
y = y.reshape(_local_dimsd.dim, axis=_local_dimsd.axis)
return y
else:
raise ValueError('expected 1-d DistributedArray, got %r'
% x.global_shape)
msg = (
"Wrong shape.\nExpects either a 1d array or, an ndarray of "
f"size `dims` when `dims` and `dimsd` both are available.\n"
f"Instead, received an array of shape {x.global_shape}."
)
raise ValueError(msg)

def adjoint(self):
"""Adjoint MPI LinearOperator
Expand Down Expand Up @@ -339,6 +359,10 @@ def __add__(self, x):
self._copy_attributes(
Op
)
if len(self.dims) == 1:
Op.dims = x.dims
if len(self.dimsd) == 1:
Op.dimsd = x.dimsd
return Op

def __neg__(self):
Expand All @@ -355,20 +379,28 @@ def _adjoint(self):
Op = _AdjointLinearOperator(self)
self._copy_attributes(
Op,
exclude=['dims', 'dimsd']
exclude=['dims', 'dimsd', '_local_dims', '_local_dimsd']
)
Op.dims = self.dimsd
Op.dimsd = self.dims
if getattr(self, "_local_dims", None) is not None:
Op._local_dimsd = self._local_dims
if getattr(self, "_local_dimsd", None) is not None:
Op._local_dims = self._local_dimsd
return Op

def _transpose(self):
Op = _TransposedLinearOperator(self)
self._copy_attributes(
Op,
exclude=['dims', 'dimsd']
exclude=['dims', 'dimsd', '_local_dims', '_local_dimsd']
)
Op.dims = self.dimsd
Op.dimsd = self.dims
if getattr(self, "_local_dims", None) is not None:
Op._local_dimsd = self._local_dims
if getattr(self, "_local_dimsd", None) is not None:
Op._local_dims = self._local_dimsd
return Op

def conj(self):
Expand All @@ -388,7 +420,7 @@ def _copy_attributes(
exclude: list[str] | None = None,
) -> None:
"""Copy attributes from one MPILinearOperator to another"""
attrs = ["dims", "dimsd"]
attrs = ["dims", "dimsd", "_local_dims", "_local_dimsd"]
if exclude is not None:
for item in exclude:
attrs.remove(item)
Expand Down
37 changes: 24 additions & 13 deletions pylops_mpi/basicoperators/BlockDiag.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
from typing import Optional, Sequence, List
from types import SimpleNamespace
import numpy as np
from scipy.sparse.linalg._interface import _get_dtype
from mpi4py import MPI
from typing import Optional, Sequence, List
from numbers import Integral

from pylops import LinearOperator
Expand Down Expand Up @@ -109,37 +110,47 @@ def __init__(self, ops: Sequence[LinearOperator],
nops[iop] = oper.shape[0]
mops[iop] = oper.shape[1]
self.mops = mops.sum()
self.local_shapes_m = base_comm.allgather((self.mops, ))
self.nops = nops.sum()
self.local_shapes_n = base_comm.allgather((self.nops, ))
dimsd = [d for dimsd in base_comm.allgather([op.dimsd for op in self.ops]) for d in dimsd]
if len(set(dimsd)) == 1:
self._local_dimsd = SimpleNamespace(dim=(len(ops), *dimsd[0]), axis=0)
dimsd = (base_comm.allreduce(len(ops)), *dimsd[0])
else:
self._local_dimsd = SimpleNamespace(dim=(self.nops, ), axis=0)
dimsd = (base_comm.allreduce(self.nops),)
dims = [d for dims in base_comm.allgather([op.dims for op in self.ops]) for d in dims]
if len(set(dims)) == 1:
self._local_dims = SimpleNamespace(dim=(len(ops), *dims[0]), axis=0)
dims = (base_comm.allreduce(len(ops)), *dims[0])
else:
self._local_dims = SimpleNamespace(dim=(self.mops, ), axis=0)
dims = (base_comm.allreduce(self.mops),)
self.local_shapes_m = base_comm.allgather((self.mops, ))
self.local_shapes_n = base_comm.allgather((self.nops,))
self.nnops = np.insert(np.cumsum(nops), 0, 0)
self.mmops = np.insert(np.cumsum(mops), 0, 0)
dimsd = (base_comm.allreduce(self.nops), )
dims = (base_comm.allreduce(self.mops), )
dtype = _get_dtype(ops) if dtype is None else np.dtype(dtype)
super().__init__(dims=dims, dimsd=dimsd, dtype=dtype, base_comm=base_comm)

@reshaped(forward=True, stacking=True)
def _matvec(self, x: DistributedArray) -> DistributedArray:
ncp = get_module(x.engine)
y = DistributedArray(global_shape=self.shape[0], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, local_shapes=self.local_shapes_n,
mask=self.mask, engine=x.engine, dtype=self.dtype)
y = DistributedArray(global_shape=self.shape[0], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl,
local_shapes=self.local_shapes_n, mask=self.mask, engine=x.engine, dtype=self.dtype)
y1 = []
for iop, oper in enumerate(self.ops):
y1.append(oper.matvec(x.local_array[self.mmops[iop]:
self.mmops[iop + 1]]))
y1.append(oper.matvec(x.local_array[self.mmops[iop]: self.mmops[iop + 1]]))
y[:] = ncp.concatenate(y1)
return y

@reshaped(forward=False, stacking=True)
def _rmatvec(self, x: DistributedArray) -> DistributedArray:
ncp = get_module(x.engine)
y = DistributedArray(global_shape=self.shape[1], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, local_shapes=self.local_shapes_m,
mask=self.mask, engine=x.engine, dtype=self.dtype)
y = DistributedArray(global_shape=self.shape[1], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl,
local_shapes=self.local_shapes_m, mask=self.mask, engine=x.engine, dtype=self.dtype)
y1 = []
for iop, oper in enumerate(self.ops):
y1.append(oper.rmatvec(x.local_array[self.nnops[iop]:
self.nnops[iop + 1]]))
y1.append(oper.rmatvec(x.local_array[self.nnops[iop]: self.nnops[iop + 1]]))
y[:] = ncp.concatenate(y1)
return y

Expand Down
4 changes: 3 additions & 1 deletion pylops_mpi/basicoperators/FirstDerivative.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from typing import Callable, Union
from types import SimpleNamespace
import numpy as np
from mpi4py import MPI

Expand All @@ -11,7 +12,7 @@
MPILinearOperator,
Partition
)

from pylops_mpi.DistributedArray import local_split
from pylops_mpi.utils.decorators import reshaped


Expand Down Expand Up @@ -90,6 +91,7 @@ def __init__(self,
base_comm: MPI.Comm = MPI.COMM_WORLD,
dtype: DTypeLike = np.float64):
dims = _value_or_sized_to_tuple(dims)
self._local_dims = self._local_dimsd = SimpleNamespace(dim=local_split(dims, base_comm, Partition.SCATTER, axis=0), axis=0)

@mrava87 mrava87 Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Mmh I have not tried if it is possible but say we define the input DistributedArray with custom local_shapes (x) and then Op=MPIFirstDerivative and then do Op @ x... this is legit right, but now the local_split dims would not match the dims of the input array?

super().__init__(dims=dims, dimsd=dims, dtype=np.dtype(dtype), base_comm=base_comm)
self.sampling = sampling
self.kind = kind
Expand Down
3 changes: 3 additions & 0 deletions pylops_mpi/basicoperators/SecondDerivative.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from typing import Callable, Union
from types import SimpleNamespace
import numpy as np
from mpi4py import MPI

Expand All @@ -8,6 +9,7 @@

from pylops_mpi import DistributedArray, MPILinearOperator, Partition
from pylops_mpi.utils.decorators import reshaped
from pylops_mpi.DistributedArray import local_split


class MPISecondDerivative(MPILinearOperator):
Expand Down Expand Up @@ -81,6 +83,7 @@ def __init__(
dtype: DTypeLike = np.float64,
) -> None:
dims = _value_or_sized_to_tuple(dims)
self._local_dims = self._local_dimsd = SimpleNamespace(dim=local_split(dims, base_comm, Partition.SCATTER, axis=0), axis=0)
super().__init__(dims=dims, dimsd=dims, dtype=np.dtype(dtype), base_comm=base_comm)
self.sampling = sampling
self.kind = kind
Expand Down
17 changes: 12 additions & 5 deletions pylops_mpi/basicoperators/VStack.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
from typing import Sequence, Optional
from types import SimpleNamespace
import numpy as np
from scipy.sparse.linalg._interface import _get_dtype
from mpi4py import MPI
from typing import Sequence, Optional

from pylops import LinearOperator
from pylops.utils import DTypeLike
Expand Down Expand Up @@ -106,15 +107,21 @@ def __init__(self, ops: Sequence[LinearOperator],
for iop, oper in enumerate(self.ops):
nops[iop] = oper.shape[0]
self.nops = nops.sum()
self.local_shapes_n = base_comm.allgather((self.nops, ))
mops = [oper.shape[1] for oper in self.ops]
mops = np.concatenate(base_comm.allgather(mops), axis=0)
if len(set(mops)) > 1:
raise ValueError("Operators have different number of columns")
self.mops = int(mops[0])
self.nnops = np.insert(np.cumsum(nops), 0, 0)
dimsd = (base_comm.allreduce(self.nops), )
dims = (self.mops, )
dims = [d for dims in base_comm.allgather([op.dims for op in self.ops]) for d in dims]
self.local_shapes_n = base_comm.allgather((self.nops,))
if len(set(dims)) == 1:
dims = dims[0]
else:
dims = (self.mops,)
self._local_dimsd = SimpleNamespace(dim=(self.nops, ), axis=0)
self._local_dims = SimpleNamespace(dim=dims, axis=0)
dtype = _get_dtype(self.ops) if dtype is None else np.dtype(dtype)
super().__init__(dims=dims, dimsd=dimsd, dtype=dtype, base_comm=base_comm)

Expand All @@ -124,8 +131,8 @@ def _matvec(self, x: DistributedArray) -> DistributedArray:
raise ValueError(f"x should have partition={Partition.BROADCAST},{Partition.UNSAFE_BROADCAST}"
f"Got {x.partition} instead...")
# the output y should use NCCL if the operand x uses it
y = DistributedArray(global_shape=self.shape[0], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, local_shapes=self.local_shapes_n,
engine=x.engine, dtype=self.dtype)
y = DistributedArray(global_shape=self.shape[0], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl,
local_shapes=self.local_shapes_n, engine=x.engine, dtype=self.dtype)
y1 = []
for iop, oper in enumerate(self.ops):
y1.append(oper.matvec(x.local_array))
Expand Down
24 changes: 15 additions & 9 deletions pylops_mpi/utils/decorators.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import numpy as np

from pylops_mpi import DistributedArray, Partition
from pylops_mpi.DistributedArray import local_split


def reshaped(

@mrava87 mrava87 Jul 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I follow this but I still have a question 😉

Before we always assumed that the input of a _matvec/_rmatvec was a 1D Distributed Array so when we decorated them with @Reshaped, inside here we will always have to create a new ND Distributed array and fill it (for a moment I am ignoring the ghost cell path).

Now, we want to enable the scenario when a ND Distributed array is passed directly to the _matvec/_rmatvec of an operator, so when we get into this the array may already by ND. At that stage if cells_front=cells_back=0, would you not say that this method could just pass x to f?

Also, again along similar lines

if len(y.global_shape) > 1:
   # Make sure y is distributed along axis=0 before applying ravel
    y = y.redistribute(axis=0).ravel()

does this not mean that the output is still again always flattened, so we can't yet have an operator that takes a ND Distributed arrays and returns a ND Distributed array? Maybe somehow related to you second point ('I actually...) But then, is this not defeating a bit the real purpose of this PR?

PS: whether we can or not enable ND Arrays flowing through ops and solvers, one by-product of this PR, which we should definitely merge is the improvement in efficiency when it comes to the setting of local_array in various places 😄

Expand Down Expand Up @@ -56,14 +57,9 @@ def wrapper(self, x: DistributedArray):
and f.__name__ != "div"
and f.__name__ != "__truediv__"
)
local_shapes = None
global_shape = getattr(self, "dims") if fwd else getattr(self, "dimsd", getattr(self, "dims"))
arr = DistributedArray(global_shape=global_shape,
base_comm=x.base_comm,
base_comm_nccl=x.base_comm_nccl,
local_shapes=local_shapes, axis=0,
engine=x.engine, dtype=x.dtype)
arr_local_shapes = np.asarray(arr.base_comm.allgather(np.prod(arr.local_shape)))
global_shape = getattr(self, "dims") if fwd else getattr(self, "dimsd")
local_shapes = self.base_comm.allgather(local_split(global_shape, self.base_comm, Partition.SCATTER, 0))
arr_local_shapes = np.array([np.prod(shape) for shape in local_shapes])
x_local_shapes = np.asarray(x.base_comm.allgather(np.prod(x.local_shape)))
# Calculate num_ghost_cells required for each rank
dif = np.cumsum(arr_local_shapes - x_local_shapes)
Expand All @@ -74,7 +70,17 @@ def wrapper(self, x: DistributedArray):
ghosted_array = x.add_ghost_cells(cells_front=cells_front, cells_back=cells_back)
# Fill the redistributed array
index = max(0, dif[self.rank - 1])
arr[:] = ghosted_array[index: arr_local_shapes[self.rank] + index].reshape(arr.local_shape)
local_array = ghosted_array[index:index + arr_local_shapes[self.rank]].reshape(local_shapes[self.rank])
arr = DistributedArray(
global_shape=global_shape,
base_comm=x.base_comm,
base_comm_nccl=x.base_comm_nccl,
local_shapes=local_shapes,
engine=x.engine,
dtype=x.dtype,
local_array=local_array,
axis=0
)
y: DistributedArray = f(self, arr)
if len(y.global_shape) > 1:
# Make sure y is distributed along axis=0 before applying ravel
Expand Down
Loading
Loading