Handle dtensor spec in sharded view

[pzharrington]

PhysicsNeMo Pull Request

Description

Updates _sharded_view_forward to handle DTensor spec as well, fixing a bug occurring in DiT backward pass from the reshape ops.

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.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Summary

This PR fixes a bug in _sharded_view_forward where a plain DTensor (not a ShardTensor subclass) could be passed during the backward pass of ShardedView. The plain DTensorSpec lacks the _sharding_shapes attribute that ShardTensorSpec provides, causing an AttributeError. The fix detects plain DTensors and reconstructs sharding shapes using compute_sharding_shapes_from_chunking_global_shape, which infers per-rank shapes from chunk semantics.

• The approach is sound: isinstance(tensor, DTensor) and not isinstance(tensor, ShardTensor) correctly identifies the case, and the chunk-based reconstruction is appropriate for autograd-produced gradients.
• The function's type annotation (tensor: ShardTensor) should be updated to reflect that it now accepts plain DTensor inputs.
• The chunk-based reconstruction assumes even/chunk-style sharding. A comment documenting this assumption would help future maintainers.
• No new tests were added for the plain DTensor code path. The PR checklist item for test coverage is unchecked.

Important Files Changed

Filename	Overview
physicsnemo/domain_parallel/shard_utils/view_ops.py	Adds DTensor spec handling in _sharded_view_forward to reconstruct sharding shapes from chunk semantics when the input is a plain DTensor (not a ShardTensor). Type annotation not updated to reflect the broader accepted input type.

_{Last reviewed commit: 4d24830}

[greptile-apps]

_{1 file reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

physicsnemo/domain_parallel/shard_utils/view_ops.py
Type annotation doesn't match actual usage

The function now accepts plain DTensor objects (not just ShardTensor) via the is_plain_dtensor check on line 457, but the type hint still declares tensor: ShardTensor. The docstring's Parameters section also describes the input as ShardTensor. Consider updating both to DTensor | ShardTensor (or just DTensor since ShardTensor is a subclass) so callers and static analysis tools understand the accepted input types.

def _sharded_view_forward(
    tensor: DTensor, target_shape: Sequence[int]
) -> ShardTensor:

[coreyjadams]

Approving to unblock. Hoping to get a real fix in ASAP.

[coreyjadams]

/blossom-ci