feat(distributed): add type annotations and docstrings to utils.py

[harshaa765]

Summary

• Addresses the TODO at line 17 of physicsnemo/distributed/utils.py requesting improved docstrings
• Adds full PEP 484 type annotations to all previously untyped functions: get_memory_format, pad_helper, truncate_helper, split_tensor_along_dim, distributed_transpose, _reduce, and _split
• Adds the missing Optional[float] return type to reduce_loss
• Expands one-liner docstrings into NumPy-style docs with Parameters, Returns, and (where applicable) Raises sections
• Adds Tuple to the typing imports to support the new return annotations
• Zero logic changes — purely annotations and documentation

Motivation

The distributed utilities are on the critical path for all multi-GPU training runs. Proper type annotations improve IDE autocomplete, enable static analysis via mypy, and make the implicit contracts (e.g. _reduce returns the same tensor object when no up-cast is needed; distributed_transpose returns None for the work handle in synchronous mode) explicit in the signature.

Test plan

• Confirm existing tests pass: pytest test/distributed/
• Optionally run mypy physicsnemo/distributed/utils.py to verify annotations are consistent

[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.

[greptile-apps]

Greptile Summary

This PR is a pure documentation and type-annotation uplift for physicsnemo/distributed/utils.py, with zero logic changes. It expands seven one-liner docstrings into full NumPy-style docs and adds PEP 484 type annotations throughout.

• Adds Parameters, Returns, and Raises sections to get_memory_format, pad_helper, truncate_helper, split_tensor_along_dim, distributed_transpose, _reduce, and _split.
• Adds the missing Optional[float] return annotation to reduce_loss and imports Tuple from typing.

Important Files Changed

Filename	Overview
physicsnemo/distributed/utils.py	Pure documentation / annotation update: adds NumPy-style docstrings and PEP 484 type hints to seven functions; three minor docstring gaps remain (TODO not removed, reduce_loss missing Returns section, split_tensor_along_dim Raises wording inaccurate).

Comments Outside Diff (1)

1. physicsnemo/distributed/utils.py, line 200-204 (link)

   The reduce_loss function now has an Optional[float] return type annotation, but the existing docstring still has no Returns section. The None return path (non-destination ranks) is implicit, so adding a Returns block here would complete the documentation that this PR is meant to improve.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "feat(distributed): add type annotations ..." | Re-trigger Greptile}

[greptile-apps]

The TODO comment the PR explicitly states it is addressing is still present. Since the docstrings have been added, the comment should be removed to keep the file tidy.

Suggested change

	# TODO this also needs more docstrings
	from typing import List, Optional
	from typing import List, Optional, Tuple
	from typing import List, Optional, Tuple

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[peterdsharpe]

Suggested change

	tensor: torch.Tensor, dim: int, new_size: int, mode: str = "zero"
	tensor: torch.Tensor, dim: int, new_size: int, mode: Literal["zero", "conj"] = "zero"

+ add from typing import Literal up top

[harshaa765]

Good suggestion — applied in 84947d1. Narrowed mode from str to Literal["zero", "conj"] and added Literal to the typing imports. Updated the docstring parameter type to match as well.