Add sparse COO/CSR/CSC/BSR/BSC meta tensor input support to torch.sum (pytorch#121673)

As in the title. Fixes an issue reported in pytorch#117907 (comment) Pull Request resolved: pytorch#121673 Approved by: https://github.com/cpuhrsch

Author: pearu

aten/src/ATen/native/ReduceOpsUtils.h

@@ -368,7 +368,13 @@ static void resize_reduction(
  DimVector dims_ = at::native::make_dim_vector(opt_dims, self.dim());
  maybe_wrap_dims(dims_, self.dim());
  auto shape = get_reduction_shape(self, dims_, keepdim, allow_empty_dims);
- meta.set_output_raw_strided(0, shape, {}, self.options().dtype(out_dtype));
+ if (self.layout() == kStrided) {
+ meta.set_output_raw_strided(0, shape, {}, self.options().dtype(out_dtype));
+ } else if (shape.size() == 0) {
+ meta.set_output_raw_strided(0, shape, {}, self.options().dtype(out_dtype).layout(kStrided));
+ } else {
+ TORCH_CHECK(false, "resize_reduction: support for output with ", self.layout(), " layout is not implemented yet");
+ }
  namedinference::propagate_names_for_reduction(
  meta.maybe_get_output(), self, dims_, keepdim);
 }

aten/src/ATen/native/native_functions.yaml

@@ -5695,8 +5695,8 @@
  variants: function, method
  dispatch:
  CompositeExplicitAutograd: sum
- SparseCPU, SparseCUDA: sum_coo
- SparseCsrCPU, SparseCsrCUDA: sum_csr
+ SparseCPU, SparseCUDA, SparseMeta: sum_coo
+ SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sum_csr
  autogen: sum.out
 
 - func: sum.dim_IntList(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor

test/test_sparse.py

@@ -4446,6 +4446,18 @@ def test_zeros_like_fake(self, dtype, layout):
  result = torch.zeros_like(f, device=f.fake_device)
  self.assertEqual(result, expected)
 
+ @all_sparse_layouts('layout', include_strided=False)
+ @parametrize("dtype", [torch.float64])
+ def test_sum_meta(self, dtype, layout):
+ device = 'cpu'
+ index_dtype = torch.int64
+ for t in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype):
+ m = t.to(device='meta')
+ r = torch.sum(m)
+ self.assertEqual(r.layout, torch.strided)
+ self.assertTrue(r.is_meta)
+ self.assertEqual(r.shape, ())
+
 
 class _SparseDataset(torch.utils.data.Dataset):
  # An utility class used in TestSparseAny.test_dataloader method.