physicsnemo-shard-tensor

Build a 2D DeviceMesh with named ddp and domain axes via DistributedManager.initialize_mesh, scatter inputs over the domain mesh with scatter_tensor using a Shard placement on a spatial dimension, keep the model code unchanged, and wrap with DistributedDataParallel on the explicit ddp process group. Per-domain-group batch size must be 1.

Register a patch from user code at import time using ShardTensor.register_function_handler, register_dispatch_handler, or register_named_function_handler, without forking physicsnemo. Model the implementation on existing patches in physicsnemo/domain_parallel/shard_utils/ such as pooling_patches.py, then validate with numerical_shard_tensor_check using check_grads=True.

Use DistributedDataParallel when all parameters are plain tensors, since auto-promotion keeps them plain even with domain parallelism. Use FSDP2 fully_shard only when you need parameter sharding for memory or have DTensor spatial parameters; never use FSDP1, which conflicts with the auto-promotion design.

The TypeError is misleading: binary dunders convert an internal NotImplementedError into NotImplemented, so CPython emits the generic message and hides the real traceback. Temporarily replace x + y with torch.add(x, y) to surface the true underlying exception.

Do not use it for generic PyTorch DDP/FSDP setup without domain parallelism, LLM tensor or pipeline parallelism such as Megatron-style sharding, single-GPU training, or choosing PhysicsNeMo models and examples. ShardTensor targets spatial and sequence sharding of activations for physics workloads.

Yes, with constraints: pass dynamic=False, and when domain_size is greater than 1, compile regionally while leaving sharded ring attention eager. Call torch._dynamo.reset() between input-size changes, and verify compiled versus eager numerics.