Good catch. Could you eliminate using AT_DISPATCH_INDEX_TYPES as index_t is not used in this block anymore?

Thanks for your comments. I find that simply changing index_t to int64_t here is not that appropriate. The root cause should be index_type should be align with columns_map's dtype, not col_indices's dtype. Have updated.

SparseCsrCPU/CUDA dispatch keys represent all sparse compressed layouts: CSR, BSR, CSC, BSC. The same holds for the sum_sparse_csr function so I suggest:

Use:

Conversion to a strided tensor is unnecessary, using self.to_sparse_csr() should be sufficient and is more efficient. That said, I think no conversions are necessary because we have the following invariants:

batch_dim = csc.dim() - csc.dense_dim() - csc.sparse_dim()
csc.layout == torch.sparse_csc
csc.transpose(batch_dim, batch_dim + 1).layout == torch.sparse_csr
torch.sum(csc, dim=dim, keepdim=True) == torch.sum(csc.transpose(batch_dim, batch_dim+1),
      dim=(*dim[:batch_dim], dim[batch_dim+1], dim[batch_dim], *dim[batch_dim+2:]),
      keepdim=True).transpose(batch_dim, batch_dim+1)

So, I suggest using the following method for computing reductions on a CSC tensor:

Nit:

Optimization nit continued:

This change is not needed as it will disable existing tests for CSR and CSC samples.

reductions over sparse dimensions of sparse compressed tensors require keepdim=True

IIUC, dim is an integer that does not correspond to a new feature that this PR implements.

If we don't implement support for non-CSR layouts, there is no need to check other parameters. So:

If not check the parameters, all the RuntimeError would be sum expected input with strided or sparse_csr layout, got layout, which could not pass the cases which expected error torch.empty: Only batched sparse compressed. WDYT?

The expected error torch.empty: Only batched sparse compressed.. is just wrong and it should be updated in the sparse.py after the error message is updated.

If we don't implement support for non-CSR layouts, there is no need to check other parameters. So:

Nit:

Nit:

Just a note: in principle, the need_acc information is defined by scalar_t and ReductionOp. This would allow compile-time optimizations because the if-blocks on need_acc can use constexpr. For instance, if ReductionOp would store its template typename as a type member, we could have:

constexpr bool need_acc = !std::is_same<scalar_t, rop::type>::value;

We should have the following invariant holding:

sparse.sum(dim=dim, keepdim=True).to_dense() == sparse.to_dense().sum(dim=dim, keepdim=True)

So:

(assertEqual with handle the conversion of sparse_sum to a strided tensor).

These tests are already covered by test_sum. So:

In fact, considering the tests test_reductions and test_reductions_backward in test_sparse.py, there is no need for test_sum_dim_reduce.

Just verify that sample_inputs_sparse_reduction_sum in torch/testing/_internal/opinfo/definitions/sparse.py produce the equivalent samples (I believe it does) and remove test_sum_dim_reduce.

Now that need_acc is constexpr, I suggest using:

Hmm, this is wrong as it overrides the user-specified dtype value. For strided tensors, we have sum result dtype tests in test/test_reductions.py but it looks like we don't generate sparse samples for sparse reduction tests (test_reductions) in test/test_sparse.py that would have caught this issue.
I think we should have:

or similar.

Same note here as above: when a user specifies dtype then the result must have the specified dtype.

Why this change? This disables the consistency checks for CSC and CSR tensors and we should not do that.

These changes are due to _validate_sample_input_sparse_reduction_sum make the limit to CSC tensor and CSR tensor with bool or complex32.
Therefore the sample will not be added into samples in test_sparse_csr.py, and raise Error of Expected at least one 2 or higher D tensor in samples.

WDYT? @pearu

samples that don't pass _validate_sample_input_sparse_reduction_sum (read: such samples are mapped to ErrorInput instances) are not included in the set of samples that op.sample_inputs(device, dtype) generates. So, this change is unnecessary (it ought to be).

In addition, notice that test_consistency is used for many operations, not just sum. So, any changes to this test must not affect testing other operations. This PR just does this: it disables consistency tests for all non-sum operations with respect to CSC inputs.

So, please undo this change here and report what problems exist for the sum operation. If there are any, these problems ought to be tackled in sparse.py, not here.

Thanks for the clarification, @pearu . I have removed the unnecessary change in test_sparse_csr.py.

Now the problem is, in sparse.py, CSC inputs and CSR inputs with bool and complex32 data type will generate ErrorInput, which is expected in this PR. And the samples will become empty, and raise Error of Expected at least one 2 or higher D tensor in samples. How can we do in sparse.py to fix the error? Thank you!

            if validate_sample_input_sparse(op, sparse_sample, check_validate=False) is not sparse_sample:
                # that is, the validation returns the sparse sample
                # wrapped within ErrorInput instance
                continue
            samples.append((sample, sparse_sample))

        # Fail early to prevent silent success with this test
        if len(samples) == 0:
            raise ValueError("Expected at least one 2 or higher D tensor in samples.")

Can you update this to "Currently the only compressed sparse format supported for sum.dim_IntList is CSR, but got layout ", layout. It's also not true that only strided and CSR are supported. We also have a kernel registered under the Sparse dispatch key which means support for COO.

Did you check whether we already have support for this conversion? We have the torch.Tensor.to operator.

Please refer to the comment #99292 (comment).

This line is repeated 4 times. The intent of my higher level comment was not around a single block of particular code. It was to abstract away repeated code.

I think we can also include the construction and resize of Tensor new_values, new_values_acc; into this helper function.

Thank you for addressing the comments. Is there something you can do to abstract away this line as well? It's repeat 5 times.