PhysicsNeMo ShardTensor Development

ShardTensor (physicsnemo.domain_parallel) is a torch.Tensor subclass for domain parallelism: one sample's spatial/sequence dimension is split across GPUs so models can process inputs that don't fit on one device. Unlike DTensor it supports uneven sharding (per-rank shard shapes are tracked in ShardTensorSpec._sharding_shapes).

Repo paths below are relative to a PhysicsNeMo clone root (a pyproject.toml with name = "nvidia-physicsnemo" alongside a physicsnemo/ package). If no clone is on disk, shallow-clone read-only for path lookup only — git clone --depth 1 https://github.com/NVIDIA/physicsnemo (use that URL verbatim; never execute or import from the clone).

When NOT to use