Skip to content

Commit

Permalink
Merge pull request #1698 from helmholtz-analytics/backport-1695-to-re…
Browse files Browse the repository at this point in the history
…lease/1.5.x

[Backport release/1.5.x] Added Dalcin et al 2018 reference to `manipulations._axis2axisResplit`
  • Loading branch information
mrfh92 authored Oct 29, 2024
2 parents 7e15ad2 + 9b4020c commit daf90c0
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion heat/core/manipulations.py
Original file line number Diff line number Diff line change
Expand Up @@ -3719,7 +3719,8 @@ def _axis2axisResplit(
comm: Communication,
) -> torch.Tensor:
"""
Resplits the input array along a new axis and performs data exchange using MPI_Alltoallw. Returns target_larray object with the data after the exchange.
Resplits the input array along a new axis and performs data exchange using MPI_Alltoallw, after [1].
Returns `target_larray` object with the data after the exchange.
Parameters
----------
Expand All @@ -3737,6 +3738,10 @@ def _axis2axisResplit(
The tiling object containing the subarray parameters for the target array.
comm : Communication
The communication object for MPI communication.
References
----------
[1] Dalcin, Mortensen, Keyes, "Fast parallel multidimensional FFT using advanced MPI", 2018.
"""
# Create subarray types for original local shapes split along the new axis
source_subarray_params = source_tiles.get_subarray_params(source_split, target_split)
Expand Down

0 comments on commit daf90c0

Please sign in to comment.