Fix missing shapes argument in DistributedIRFFT2 gather

[Anai-Guo]

PhysicsNeMo Pull Request

Description

conj_pad_helper_2d in physicsnemo/distributed/fft.py calls
gather_from_parallel_region without the required shapes argument:

tensor_pad_gather = gather_from_parallel_region(
    tensor_pad, dim=other_dim, group="spatial_parallel"
)

but the helper is defined in physicsnemo/distributed/mappings.py as

def gather_from_parallel_region(input, dim, shapes, group):

shapes is positional-and-required and there is no **kwargs, so this raises

TypeError: gather_from_parallel_region() missing 1 required positional argument: 'shapes'

conj_pad_helper_2d is reached from DistributedIRFFT2.forward (fft.py:191),
so every DistributedIRFFT2 forward pass fails. The failure happens at
argument-binding time, before any collective is issued, so it is not
GPU/rank-count dependent.

This is the only one of the five gather_from_parallel_region call sites that
omits shapes — the other four all pass it:

Call site	Passes shapes?
models/afno/distributed/afno.py:201	✅ shapes=scatter_shapes
models/afno/distributed/layers.py:226	✅ shapes=self.gather_shapes
models/afno/distributed/layers.py:353	✅ shapes=self.in_shapes
models/afno/distributed/layers.py:631	✅ shapes=gather_shapes
distributed/fft.py:41	❌ missing

Why shapes=None

shapes=None is explicitly supported and means "assume equal shape on all
ranks" (all_gather_v_wrapper in distributed/utils.py, whose own sizes
parameter defaults to None). That is the right value here:

• Both DistributedRFFT2 and DistributedIRFFT2 document their contract as
  "a single global tensor which is distributed along a specified dimension
  into chunks of equal size".
• The gather is immediately paired with
  scatter_to_parallel_region(tensor_pad_gather, dim=other_dim, group="spatial_parallel")
  three statements later, which takes no shapes and splits evenly — so an
  even gather is what keeps the round trip symmetric.

The change is therefore behaviour-preserving relative to the documented intent;
it only makes the call actually bind.

Verification

Since these paths are # pragma: no cover and need a multi-rank job, I verified
by replaying the real signature against the real call site with
inspect.Signature.bind (AST-extracted from the unmodified sources, so no
hand-copying):

BEFORE line 41: TypeError: missing a required argument: 'shapes'
AFTER  line 41: BINDS OK

ruff check / ruff format --check (v0.12.5, the pinned pre-commit version)
pass on the changed file.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Note on tests: the distributed FFT primitives have no existing multi-rank test
and are marked # pragma: no cover, so I did not add one here. Happy to add a
torchrun-based test if you would like it in this PR.

Dependencies

None.

---

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

CODEOWNERS review map

Current for commit c8cac711af38. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @coreyjadams — 1 file(s)

• physicsnemo/distributed/fft.py

⏳ @negin513 — 1 file(s)

• physicsnemo/distributed/fft.py

No CODEOWNER

• CHANGELOG.md

Comment /codeowners-info to refresh.

[greptile-apps]

Greptile Summary

Adds the required shapes argument to the distributed inverse FFT gather, preventing argument binding from failing before communication begins.

• Passes shapes=None to preserve the transform’s documented equal-shard behavior.
• Records the DistributedIRFFT2 fix in the changelog.

Important Files Changed

Filename	Overview
physicsnemo/distributed/fft.py	Correctly supplies the required gather argument while retaining the documented equal-size shard assumption.
CHANGELOG.md	Accurately documents the repaired DistributedIRFFT2 argument-binding failure.

_{Reviews (1): Last reviewed commit: "fix: pass missing `shapes` argument in D..." | Re-trigger Greptile}