mesh: enable ShardTensor support for mesh conversion/geometry paths

physicsnemo/domain_parallel/__init__.py

@@ -44,19 +44,17 @@
 
 
 if ST_AVAILABLE:
-    # In minumum versions are met, we can import the shard tensor and spec.
+    # If minimum versions are met, we can import the shard tensor and spec.
+    # Import tensor op handlers here because they are valid on CPU and CUDA.
 
     from ._shard_tensor_spec import ShardTensorSpec
+    from .custom_ops import _tensor_ops  # noqa: F401
     from .shard_tensor import ShardTensor, scatter_tensor
 
     def register_custom_ops():
         # These imports will register the custom ops with the ShardTensor class.
         # It's done here to avoid an import cycle.
-        from .custom_ops import (  # noqa: F401
-            _tensor_ops,
-            mean_wrapper,
-            sum_wrapper,
-        )
+        from .custom_ops import mean_wrapper, sum_wrapper  # noqa: F401
         from .shard_utils import register_shard_wrappers
 
         register_shard_wrappers()

physicsnemo/domain_parallel/custom_ops/_tensor_ops.py

@@ -32,15 +32,220 @@
     Shard,
 )
 
-from physicsnemo.domain_parallel import ShardTensor
 from physicsnemo.domain_parallel._shard_tensor_spec import (
     ShardTensorSpec,
     _stride_from_contiguous_shape_C_style,
 )
+from physicsnemo.domain_parallel.shard_tensor import ShardTensor
 
 aten = torch.ops.aten
 
 
+def _extract_cross_inputs(
+    args: tuple[Any, ...], kwargs: dict[str, Any]
+) -> tuple[ShardTensor, ShardTensor]:
+    r"""Extract and validate ShardTensor inputs for cross-product handlers."""
+    input_tensor = args[0] if len(args) > 0 else kwargs.get("input")
+    other_tensor = args[1] if len(args) > 1 else kwargs.get("other")
+
+    _validate_cross_inputs(input_tensor, other_tensor)
+    return input_tensor, other_tensor
+
+
+def _validate_cross_inputs(input_tensor: Any, other_tensor: Any) -> None:
+    r"""Validate ShardTensor inputs for cross-product handlers."""
+    if not isinstance(input_tensor, ShardTensor) or not isinstance(
+        other_tensor, ShardTensor
+    ):
+        raise RuntimeError(
+            "cross with ShardTensor inputs requires both arguments to be ShardTensor."
+        )
+
+    if input_tensor._spec.mesh != other_tensor._spec.mesh:
+        raise RuntimeError(
+            "cross requires both ShardTensor inputs to share the same device mesh."
+        )
+    if input_tensor._spec.placements != other_tensor._spec.placements:
+        raise RuntimeError(
+            "cross requires both ShardTensor inputs to have identical placements."
+        )
+
+
+def _cross_result_from_local(
+    input_tensor: ShardTensor,
+    local_input: torch.Tensor,
+    local_result: torch.Tensor,
+) -> ShardTensor:
+    r"""Wrap a local cross-product result with matching ShardTensor metadata."""
+    sharding_shapes = (
+        input_tensor._spec.sharding_shapes()
+        if local_result.shape == local_input.shape
+        else "infer"
+    )
+    return ShardTensor.from_local(
+        local_result,
+        input_tensor._spec.mesh,
+        input_tensor._spec.placements,
+        sharding_shapes=sharding_shapes,
+    )
+
+
+def _normalize_cross_dim(
+    input_tensor: ShardTensor,
+    dim: int | None,
+    *,
+    allow_none: bool,
+    op_name: str,
+) -> int:
+    r"""Return the global cross-product dimension and reject invalid vector dims."""
+    ndim = input_tensor.ndim
+    if dim is None:
+        if not allow_none:
+            raise TypeError(f"{op_name}(): argument 'dim' must be int, not NoneType")
+        try:
+            normalized_dim = next(
+                i for i, size in enumerate(input_tensor.shape) if size == 3
+            )
+        except StopIteration:
+            raise RuntimeError(
+                "torch.cross with ShardTensor dim=None requires an input dimension "
+                "of size 3."
+            )
+    else:
+        if not isinstance(dim, int):
+            raise TypeError(
+                f"{op_name}(): argument 'dim' must be int, not {type(dim).__name__}"
+            )
+        if dim < -ndim or dim >= ndim:
+            raise IndexError(
+                "Dimension out of range "
+                f"(expected to be in range of [{-ndim}, {ndim - 1}], but got {dim})"
+            )
+        normalized_dim = dim if dim >= 0 else dim + ndim
+
+    if any(
+        isinstance(placement, Shard) and placement.dim == normalized_dim
+        for placement in input_tensor._spec.placements
+    ):
+        raise RuntimeError(
+            "cross with ShardTensor inputs is not supported along a sharded dimension."
+        )
+    return normalized_dim
+
+
+def linalg_cross_wrapper(
+    func: Callable,
+    types: tuple[Any, ...],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> ShardTensor:
+    r"""Functional wrapper for ``torch.linalg.cross`` on ShardTensor inputs."""
+    if kwargs is None:
+        kwargs = {}
+
+    if kwargs.get("out", None) is not None:
+        raise RuntimeError(
+            "torch.linalg.cross(out=...) is not supported for ShardTensor."
+        )
+
+    input_tensor, other_tensor = _extract_cross_inputs(args, kwargs)
+    dim = kwargs.get("dim", -1)
+    if len(args) > 2:
+        dim = args[2]
+    dim = _normalize_cross_dim(
+        input_tensor,
+        dim,
+        allow_none=False,
+        op_name="linalg_cross",
+    )
+
+    local_input = input_tensor.to_local()
+    local_result = torch.linalg.cross(
+        local_input,
+        other_tensor.to_local(),
+        dim=dim,
+    )
+    return _cross_result_from_local(input_tensor, local_input, local_result)
+
+
+def _linalg_cross_dispatch(
+    input_tensor: ShardTensor,
+    other_tensor: ShardTensor,
+    *,
+    dim: int = -1,
+) -> ShardTensor:
+    r"""ATen dispatch handler for ``aten.linalg_cross.default``."""
+    _validate_cross_inputs(input_tensor, other_tensor)
+    dim = _normalize_cross_dim(
+        input_tensor,
+        dim,
+        allow_none=False,
+        op_name="linalg_cross",
+    )
+    local_input = input_tensor.to_local()
+    local_result = torch.linalg.cross(
+        local_input,
+        other_tensor.to_local(),
+        dim=dim,
+    )
+    return _cross_result_from_local(input_tensor, local_input, local_result)
+
+
+def cross_wrapper(
+    func: Callable,
+    types: tuple[Any, ...],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> ShardTensor:
+    r"""Functional wrapper for ``torch.cross`` on ShardTensor inputs."""
+    if kwargs is None:
+        kwargs = {}
+
+    if kwargs.get("out", None) is not None:
+        raise RuntimeError("torch.cross(out=...) is not supported for ShardTensor.")
+
+    input_tensor, other_tensor = _extract_cross_inputs(args, kwargs)
+    dim = kwargs.get("dim", None)
+    if len(args) > 2:
+        dim = args[2]
+    dim = _normalize_cross_dim(
+        input_tensor,
+        dim,
+        allow_none=True,
+        op_name="cross",
+    )
+
+    local_input = input_tensor.to_local()
+    local_result = torch.cross(
+        local_input,
+        other_tensor.to_local(),
+        dim=dim,
+    )
+    return _cross_result_from_local(input_tensor, local_input, local_result)
+
+
+def _cross_dispatch(
+    input_tensor: ShardTensor,
+    other_tensor: ShardTensor,
+    dim: int | None = None,
+) -> ShardTensor:
+    r"""ATen dispatch handler for ``aten.cross.default``."""
+    _validate_cross_inputs(input_tensor, other_tensor)
+    dim = _normalize_cross_dim(
+        input_tensor,
+        dim,
+        allow_none=True,
+        op_name="cross",
+    )
+    local_input = input_tensor.to_local()
+    local_result = torch.cross(
+        local_input,
+        other_tensor.to_local(),
+        dim=dim,
+    )
+    return _cross_result_from_local(input_tensor, local_input, local_result)
+
+
 def _unbind_output_metadata(
     input_spec: ShardTensorSpec, dim: int
 ) -> tuple[int, list, dict[int, list[torch.Size]]]:
@@ -210,8 +415,15 @@ def unbind_wrapper(
 
 
 # Python-level function handlers (__torch_function__).
+ShardTensor.register_function_handler(torch.linalg.cross, linalg_cross_wrapper)
+ShardTensor.register_function_handler(aten.linalg_cross.default, linalg_cross_wrapper)
+ShardTensor.register_function_handler(torch.cross, cross_wrapper)
+ShardTensor.register_function_handler(torch.Tensor.cross, cross_wrapper)
+ShardTensor.register_function_handler(aten.cross.default, cross_wrapper)
 ShardTensor.register_function_handler(torch.unbind, unbind_wrapper)
 ShardTensor.register_function_handler(torch.Tensor.unbind, unbind_wrapper)
 
 # ATen-level dispatch handler (__torch_dispatch__).
+ShardTensor.register_dispatch_handler(aten.linalg_cross.default, _linalg_cross_dispatch)
+ShardTensor.register_dispatch_handler(aten.cross.default, _cross_dispatch)
 ShardTensor.register_dispatch_handler(aten.unbind.int, _unbind_dispatch)

physicsnemo/mesh/mesh.py

@@ -1461,15 +1461,15 @@ def cell_data_to_point_data(self, overwrite_keys: bool = False) -> "Mesh":
         # Shape: (n_cells * n_vertices_per_cell,)
         point_indices = self.cells.flatten()
 
-        # Corresponding cell index for each point
-        # Shape: (n_cells * n_vertices_per_cell,)
-        cell_indices = torch.arange(
-            self.n_cells, device=self.points.device
-        ).repeat_interleave(n_vertices_per_cell)
+        # Repeat each cell value once per incident vertex. This avoids mixing a
+        # ShardTensor data field with a generated dense cell-index tensor.
 
         converted = self.cell_data.apply(
             lambda cell_values: scatter_aggregate(
-                src_data=cell_values[cell_indices],
+                src_data=cell_values.unsqueeze(1)
+                .expand(-1, n_vertices_per_cell, *cell_values.shape[1:])
+                .clone()
+                .reshape(-1, *cell_values.shape[1:]),
                 src_to_dst_mapping=point_indices,
                 n_dst=self.n_points,
                 weights=None,

physicsnemo/mesh/utilities/_scatter_ops.py

@@ -82,7 +82,6 @@ def scatter_aggregate(
     >>> result = scatter_aggregate(src_data, src_to_dst, n_dst=2)
     >>> # result = [[1.5], [3.0]]  # point 0 gets mean of cells 0,1
     """
-    device = src_data.device
     dtype = src_data.dtype
 
     ### Get data shape beyond the first dimension
@@ -93,7 +92,7 @@ def scatter_aggregate(
 
     ### Fast path: unweighted sum is a single scatter_add_ with no extra work
     if weights is None and aggregation == "sum":
-        aggregated_data = torch.zeros((n_dst, *data_shape), dtype=dtype, device=device)
+        aggregated_data = src_data.new_zeros((n_dst, *data_shape), dtype=dtype)
         expanded_indices = src_to_dst_mapping.view(
             -1, *([1] * len(data_shape))
         ).expand_as(src_data)
@@ -102,7 +101,7 @@ def scatter_aggregate(
 
     ### Initialize weights if not provided
     if weights is None:
-        weights = torch.ones(len(src_to_dst_mapping), dtype=dtype, device=device)
+        weights = torch.ones_like(src_to_dst_mapping, dtype=dtype)
 
     ### Ensure weights have same dtype as data (avoid dtype mismatch in multiplication)
     if weights.dtype != dtype:
@@ -114,11 +113,7 @@ def scatter_aggregate(
     weighted_data = src_data * weights.view(weight_shape)
 
     ### Scatter-add weighted data to destinations
-    aggregated_data = torch.zeros(
-        (n_dst, *data_shape),
-        dtype=dtype,
-        device=device,
-    )
+    aggregated_data = src_data.new_zeros((n_dst, *data_shape), dtype=dtype)
 
     # Expand src_to_dst_mapping to match data dimensions
     expanded_indices = src_to_dst_mapping.view(-1, *([1] * len(data_shape))).expand_as(
@@ -134,7 +129,7 @@ def scatter_aggregate(
     ### Normalize weighted sum to weighted mean
     if aggregation == "mean":
         ### Compute sum of weights at each destination
-        weight_sums = torch.zeros(n_dst, dtype=dtype, device=device)
+        weight_sums = src_data.new_zeros((n_dst,), dtype=dtype)
         weight_sums.scatter_add_(
             dim=0,
             index=src_to_dst_mapping,