Repository: nod-ai/shark-ai
Pull request #1308: Add ops.mean for SplitPrimitiveTensor
Author: login: Alex-Vasile
Created: 2025-04-25T15:16:57Z | Merged: 2025-04-25T20:31:46Z

Only strange thing about this op that I haven't seen in others is that doing the mean across the `shard_dim` requires us to return a `ReplicatedTensor` rather than a `SplitPrimitiveTensor`, especially when `keepdim == False` since the `shard_dim` stops existing. 


Diff:
diff --git a/sharktank/sharktank/ops/sharded_impls.py b/sharktank/sharktank/ops/sharded_impls.py
index 24afcf93f4b..d1e3689b7a0 100644
--- a/sharktank/sharktank/ops/sharded_impls.py
+++ b/sharktank/sharktank/ops/sharded_impls.py
@@ -1051,11 +1051,49 @@ def mean_replicated(
     keepdim: bool,
     *,
     dtype: torch.dtype,
-) -> None:
+) -> ReplicatedTensor:
     shards = [mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards]
     return ReplicatedTensor(ts=shards)
 
 
+@mean.override(SplitPrimitiveTensor)
+def mean_split(
+    x: SplitPrimitiveTensor,
+    dim: Union[int, List[int]],
+    keepdim: bool,
+    *,
+    dtype: torch.dtype,
+) -> SplitPrimitiveTensor | ReplicatedTensor:
+    if not isinstance(dim, (list, tuple)):
+        dim = [dim]
+    dim = [d + len(x.shape) if d < 0 else d for d in dim]
+
+    if x.shard_dim not in dim:
+        # If keepdim == False and any entry in dim is smaller than shard_dim
+        # we need to offset shard_dim_new to have it point to the same dimension.
+        num_smaller_dims = sum(d < x.shard_dim for d in dim)
+        shard_dim_new = x.shard_dim - (not keepdim) * num_smaller_dims
+
+        shards = [
+            mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards
+        ]
+        return SplitPrimitiveTensor(ts=shards, shard_dim=shard_dim_new)
+    else:
+        gathered = cat(
+            [
+                (
+                    transfer_to_logical_device(shard, x.devices[0])
+                    if i != 0
+                    else barrier_on_logical_device(shard, x.devices[0])
+                )
+                for i, shard in enumerate(x.shards)
+            ],
+            dim=x.shard_dim,
+        )
+        meaned = mean(gathered, dim=dim, keepdim=keepdim, dtype=dtype)
+        return ReplicatedTensor(ts=meaned, shard_count=x.shard_count, devices=x.devices)
+
+
 @module_register_buffer.override(torch.nn.Module, ShardedTensor)
 def module_register_buffer_sharded(
     module: torch.nn.Module, name: str, tensor: ShardedTensor
diff --git a/sharktank/tests/ops/sharded_test.py b/sharktank/tests/ops/sharded_test.py
index 934f7de0f8c..38cd07709d3 100644
--- a/sharktank/tests/ops/sharded_test.py
+++ b/sharktank/tests/ops/sharded_test.py
@@ -5,7 +5,8 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import unittest
-from parameterized import parameterized
+import itertools
+from parameterized import parameterized, parameterized_class
 
 import torch
 
@@ -887,6 +888,67 @@ def testReplicatedLhsAndRhs(self):
             assert ops.equal(shard, unsharded_result)
 
 
+@parameterized_class(
+    ("keepdim", "mean_dim_delta"), list(itertools.product([True, False], [-1, 0, +1]))
+)
+class MeanTest(unittest.TestCase):
+    def setUp(self):
+        self.shape = (2, 4, 6, 8, 10)
+        self.shard_dim = 2
+        self.mean_dim = self.shard_dim + self.mean_dim_delta
+        self.mean_dims_multi = tuple(self.mean_dim + i for i in [-1, 0, +1])
+        self.shard_count = 2
+        torch.random.manual_seed(sum(self.mean_dims_multi) + 13 * self.keepdim)
+
+    def testMeanReplicated(self):
+        tensor = torch.rand(self.shape, dtype=torch.float32)
+        expected_result = ops.mean(tensor, dim=self.mean_dim, keepdim=self.keepdim)
+        actual_result = ops.mean(
+            ops.replicate(tensor, count=self.shard_count),
+            dim=self.mean_dim,
+            keepdim=self.keepdim,
+        )
+        torch.testing.assert_close(expected_result, ops.unbox_tensor(actual_result))
+
+    def testMeanSplit(self):
+        tensor = torch.rand(self.shape, dtype=torch.float32)
+        expected_result = ops.mean(tensor, dim=self.mean_dim, keepdim=self.keepdim)
+        sharded_tensor = ops.reshard_split(
+            tensor, dim=self.shard_dim, count=self.shard_count
+        )
+        actual_result = ops.mean(
+            sharded_tensor, dim=self.mean_dim, keepdim=self.keepdim
+        )
+        torch.testing.assert_close(expected_result, ops.unbox_tensor(actual_result))
+
+    def testMeanSplitNegativeDims(self):
+        if self.mean_dim_delta != -1:
+            self.skipTest(
+                "Using a specifc negative dim, so only running for different versions of 'keepdim'."
+            )
+        mean_dim = [-5, -1, -4]
+        tensor = torch.rand(self.shape, dtype=torch.float32)
+        expected_result = ops.mean(tensor, dim=mean_dim, keepdim=self.keepdim)
+        sharded_tensor = ops.reshard_split(
+            tensor, dim=self.shard_dim, count=self.shard_count
+        )
+        actual_result = ops.mean(sharded_tensor, dim=mean_dim, keepdim=self.keepdim)
+        torch.testing.assert_close(expected_result, ops.unbox_tensor(actual_result))
+
+    def testMeanSplitMultiDim(self):
+        tensor = torch.rand(self.shape, dtype=torch.float32)
+        expected_result = ops.mean(
+            tensor, dim=self.mean_dims_multi, keepdim=self.keepdim
+        )
+        sharded_tensor = ops.reshard_split(
+            tensor, dim=self.shard_dim, count=self.shard_count
+        )
+        actual_result = ops.mean(
+            sharded_tensor, dim=self.mean_dims_multi, keepdim=self.keepdim
+        )
+        torch.testing.assert_close(expected_result, ops.unbox_tensor(actual_result))
+
+
 class ReplicateTest(unittest.TestCase):
     def testReplicateReplicated(self):
         tensor = torch.rand(4, 5, dtype=torch.float32)


Reviews:
user: login: sogartar
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:19:47Z
user: login: sogartar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:22:51Z
user: login: Alex-Vasile
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:39:10Z
user: login: Alex-Vasile
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:41:36Z
user: login: sogartar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T19:52:08Z

Review comments:
diff_hunk: @@ -887,6 +888,38 @@ def testReplicatedLhsAndRhs(self):
             assert ops.equal(shard, unsharded_result)
 
 
+@parameterized_class(
+    ("keepdim", "mean_dim_delta"), list(itertools.product([True, False], [-1, 0, +1]))
+)
+class MeanTest(unittest.TestCase):
path: sharktank/tests/ops/sharded_test.py
user: login: sogartar
body: Thank you for the test.
Unfortunately, a lot of API ops, sharded or unsharded, don't have tests.
created_at: 2025-04-25T16:44:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1051,11 +1051,46 @@ def mean_replicated(
     keepdim: bool,
     *,
     dtype: torch.dtype,
-) -> None:
+) -> ReplicatedTensor:
     shards = [mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards]
     return ReplicatedTensor(ts=shards)
 
 
+@mean.override(SplitPrimitiveTensor)
+def mean_split(
+    x: SplitPrimitiveTensor,
+    dim: Union[int, List[int]],
+    keepdim: bool,
+    *,
+    dtype: torch.dtype,
+) -> SplitPrimitiveTensor | ReplicatedTensor:
+    if dim != x.shard_dim:
path: sharktank/sharktank/ops/sharded_impls.py
user: login: sogartar
body: We should raise an error for the unhandled case of when dim is not an int but a tuple. This will help us avoid nasty surprises in the future.
created_at: 2025-04-25T16:50:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -887,6 +888,38 @@ def testReplicatedLhsAndRhs(self):
             assert ops.equal(shard, unsharded_result)
 
 
+@parameterized_class(
+    ("keepdim", "mean_dim_delta"), list(itertools.product([True, False], [-1, 0, +1]))
+)
+class MeanTest(unittest.TestCase):
+    def setUp(self):
+        self.shape = (2, 4, 6)
+        self.shard_dim = 1
+        self.mean_dim = self.shard_dim + self.mean_dim_delta
+        self.shard_count = 2
+
+    def testMeanReplicated(self):
+        tensor = torch.rand(self.shape, dtype=torch.float32)
path: sharktank/tests/ops/sharded_test.py
user: login: sogartar
body: I know other tests are also guilty of this, but this will introduce nondeterminism in the test as the torch RNG was not manually seeded.
created_at: 2025-04-25T17:10:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -887,6 +888,38 @@ def testReplicatedLhsAndRhs(self):
             assert ops.equal(shard, unsharded_result)
 
 
+@parameterized_class(
+    ("keepdim", "mean_dim_delta"), list(itertools.product([True, False], [-1, 0, +1]))
+)
+class MeanTest(unittest.TestCase):
+    def setUp(self):
+        self.shape = (2, 4, 6)
+        self.shard_dim = 1
+        self.mean_dim = self.shard_dim + self.mean_dim_delta
+        self.shard_count = 2
+
+    def testMeanReplicated(self):
+        tensor = torch.rand(self.shape, dtype=torch.float32)
+        expected_result = ops.mean(tensor, dim=self.mean_dim, keepdim=self.keepdim)
+        actual_result = ops.mean(
+            ops.replicate(tensor, count=self.shard_count),
+            dim=self.mean_dim,
+            keepdim=self.keepdim,
+        )
+        assert ops.equal(expected_result, actual_result)
path: sharktank/tests/ops/sharded_test.py
user: login: sogartar
body: If for some reason we ever change the order of additions compared to the reference this will fail due to rounding loss of precision. It is better to do `torch.testing.assert_close` to allow for some error tolerance.
The same holds for the other test.
created_at: 2025-04-25T17:12:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1051,11 +1051,46 @@ def mean_replicated(
     keepdim: bool,
     *,
     dtype: torch.dtype,
-) -> None:
+) -> ReplicatedTensor:
     shards = [mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards]
     return ReplicatedTensor(ts=shards)
 
 
+@mean.override(SplitPrimitiveTensor)
+def mean_split(
+    x: SplitPrimitiveTensor,
+    dim: Union[int, List[int]],
+    keepdim: bool,
+    *,
+    dtype: torch.dtype,
+) -> SplitPrimitiveTensor | ReplicatedTensor:
+    if dim != x.shard_dim:
+        # If keepdim == False dim is smaller than shard_dim we need to offset
+        # shard_dim_new by -1 to have it point to the same dimension.
+        shard_dim_new = x.shard_dim
+        if (not keepdim) and dim < x.shard_dim:
+            shard_dim_new -= 1
+
+        shards = [
+            mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards
+        ]
+        return SplitPrimitiveTensor(ts=shards, shard_dim=shard_dim_new)
+    else:
+        gathered = cat(
+            [
+                (
+                    transfer_to_logical_device(shard, x.devices[0])
+                    if i != 0
+                    else barrier_on_logical_device(shard, x.devices[0])
+                )
+                for i, shard in enumerate(x.shards)
+            ],
+            dim=dim,
+        )
+        meaned = mean(gathered, dim=dim, keepdim=keepdim, dtype=dtype)
+        return ReplicatedTensor(ts=meaned, shard_count=x.shard_count, devices=x.devices)
path: sharktank/sharktank/ops/sharded_impls.py
user: login: sogartar
body: It is not critical for correctness, but for performance reasons I think we could change the implementation to
1. Each device reduces along dim.
2. All-reduce across the devices.
3. Divide by the total number of elements in the above reductions.

For most cases this should bring down communication a lot.
created_at: 2025-04-25T17:16:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1051,11 +1051,46 @@ def mean_replicated(
     keepdim: bool,
     *,
     dtype: torch.dtype,
-) -> None:
+) -> ReplicatedTensor:
     shards = [mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards]
     return ReplicatedTensor(ts=shards)
 
 
+@mean.override(SplitPrimitiveTensor)
+def mean_split(
+    x: SplitPrimitiveTensor,
+    dim: Union[int, List[int]],
+    keepdim: bool,
+    *,
+    dtype: torch.dtype,
+) -> SplitPrimitiveTensor | ReplicatedTensor:
+    if dim != x.shard_dim:
+        # If keepdim == False dim is smaller than shard_dim we need to offset
+        # shard_dim_new by -1 to have it point to the same dimension.
+        shard_dim_new = x.shard_dim
+        if (not keepdim) and dim < x.shard_dim:
+            shard_dim_new -= 1
+
+        shards = [
+            mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards
+        ]
+        return SplitPrimitiveTensor(ts=shards, shard_dim=shard_dim_new)
+    else:
+        gathered = cat(
+            [
+                (
+                    transfer_to_logical_device(shard, x.devices[0])
+                    if i != 0
+                    else barrier_on_logical_device(shard, x.devices[0])
+                )
+                for i, shard in enumerate(x.shards)
+            ],
+            dim=dim,
+        )
+        meaned = mean(gathered, dim=dim, keepdim=keepdim, dtype=dtype)
+        return ReplicatedTensor(ts=meaned, shard_count=x.shard_count, devices=x.devices)
path: sharktank/sharktank/ops/sharded_impls.py
user: login: sogartar
body: This can be in a later PR.
created_at: 2025-04-25T17:22:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1051,11 +1051,46 @@ def mean_replicated(
     keepdim: bool,
     *,
     dtype: torch.dtype,
-) -> None:
+) -> ReplicatedTensor:
     shards = [mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards]
     return ReplicatedTensor(ts=shards)
 
 
+@mean.override(SplitPrimitiveTensor)
+def mean_split(
+    x: SplitPrimitiveTensor,
+    dim: Union[int, List[int]],
+    keepdim: bool,
+    *,
+    dtype: torch.dtype,
+) -> SplitPrimitiveTensor | ReplicatedTensor:
+    if dim != x.shard_dim:
+        # If keepdim == False dim is smaller than shard_dim we need to offset
+        # shard_dim_new by -1 to have it point to the same dimension.
+        shard_dim_new = x.shard_dim
+        if (not keepdim) and dim < x.shard_dim:
+            shard_dim_new -= 1
+
+        shards = [
+            mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards
+        ]
+        return SplitPrimitiveTensor(ts=shards, shard_dim=shard_dim_new)
+    else:
+        gathered = cat(
+            [
+                (
+                    transfer_to_logical_device(shard, x.devices[0])
+                    if i != 0
+                    else barrier_on_logical_device(shard, x.devices[0])
+                )
+                for i, shard in enumerate(x.shards)
+            ],
+            dim=dim,
+        )
+        meaned = mean(gathered, dim=dim, keepdim=keepdim, dtype=dtype)
+        return ReplicatedTensor(ts=meaned, shard_count=x.shard_count, devices=x.devices)
path: sharktank/sharktank/ops/sharded_impls.py
user: login: Alex-Vasile
body: I considered this, my worry was about potential under/over flow. And that underlying floating point operations would no longer be the same leading to slight differences between the results.
created_at: 2025-04-25T17:39:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1051,11 +1051,46 @@ def mean_replicated(
     keepdim: bool,
     *,
     dtype: torch.dtype,
-) -> None:
+) -> ReplicatedTensor:
     shards = [mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards]
     return ReplicatedTensor(ts=shards)
 
 
+@mean.override(SplitPrimitiveTensor)
+def mean_split(
+    x: SplitPrimitiveTensor,
+    dim: Union[int, List[int]],
+    keepdim: bool,
+    *,
+    dtype: torch.dtype,
+) -> SplitPrimitiveTensor | ReplicatedTensor:
+    if dim != x.shard_dim:
path: sharktank/sharktank/ops/sharded_impls.py
user: login: Alex-Vasile
body: Darn. I missed that. I will see about implementing that.
created_at: 2025-04-25T17:41:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1051,11 +1051,46 @@ def mean_replicated(
     keepdim: bool,
     *,
     dtype: torch.dtype,
-) -> None:
+) -> ReplicatedTensor:
     shards = [mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards]
     return ReplicatedTensor(ts=shards)
 
 
+@mean.override(SplitPrimitiveTensor)
+def mean_split(
+    x: SplitPrimitiveTensor,
+    dim: Union[int, List[int]],
+    keepdim: bool,
+    *,
+    dtype: torch.dtype,
+) -> SplitPrimitiveTensor | ReplicatedTensor:
+    if dim != x.shard_dim:
+        # If keepdim == False dim is smaller than shard_dim we need to offset
+        # shard_dim_new by -1 to have it point to the same dimension.
+        shard_dim_new = x.shard_dim
+        if (not keepdim) and dim < x.shard_dim:
+            shard_dim_new -= 1
+
+        shards = [
+            mean(shard, dim=dim, keepdim=keepdim, dtype=dtype) for shard in x.shards
+        ]
+        return SplitPrimitiveTensor(ts=shards, shard_dim=shard_dim_new)
+    else:
+        gathered = cat(
+            [
+                (
+                    transfer_to_logical_device(shard, x.devices[0])
+                    if i != 0
+                    else barrier_on_logical_device(shard, x.devices[0])
+                )
+                for i, shard in enumerate(x.shards)
+            ],
+            dim=dim,
+        )
+        meaned = mean(gathered, dim=dim, keepdim=keepdim, dtype=dtype)
+        return ReplicatedTensor(ts=meaned, shard_count=x.shard_count, devices=x.devices)
path: sharktank/sharktank/ops/sharded_impls.py
user: login: sogartar
body: The operation does not guarantee in what order are additions performed. The underlying implementation may tile on a reduction dimension even on a GPU. Actually, tiling would yield better numerical stability. The accumulator would not get that big and lose precision when adding small and big numbers together.
created_at: 2025-04-25T19:52:08Z
author_association: CONTRIBUTOR

PR comments:
user: login: sogartar
created_at: 2025-04-25T16:42:59Z
body: > Only strange thing about this op that I haven't seen in others is that doing the mean across the `shard_dim` requires us to return a `ReplicatedTensor` rather than a `SplitPrimitiveTensor`, especially when `keepdim == False` since the `shard_dim` stops existing.

I think the sensible thing is to return a replicated tensor when a reduction dimension is sharded.
author_association: CONTRIBUTOR
user: login: Alex-Vasile
created_at: 2025-04-25T16:46:49Z
body: > > Only strange thing about this op that I haven't seen in others is that doing the mean across the `shard_dim` requires us to return a `ReplicatedTensor` rather than a `SplitPrimitiveTensor`, especially when `keepdim == False` since the `shard_dim` stops existing.
> 
> I think the sensible thing is to return a replicated tensor when a reduction dimension is sharded.

I agree, nothing else makes sense. I just haven't seen another op when this was required. That's what I've implemented here.
author_association: CONTRIBUTOR