[Bug] Incorrect Sparse Operator Construction from COO to CSR

[alireza78h]

🐛 Bug

The construction of the sparse operator during the conversion from COO (Coordinate format) to CSR (Compressed Sparse Row format) produces incorrect results.

To reproduce

from cola.ops import Sparse
import torch
row = torch.tensor([1, 3, 2, 2 ])
col = torch.tensor([1, 3, 2, 1 ])
data = torch.tensor([0.2643, 0.3951, 0.8338, 0.2440], dtype=torch.float64)
indices = torch.stack([row, col])
coo_tensor_torch = torch.sparse_coo_tensor(indices, data, size=[4, 4])
coo_op_cola = Sparse(data, row, col, (4,4))
print(torch.isclose(coo_op_cola.to_dense(),coo_tensor_torch.to_dense()).all())
# This outputs False
print(coo_op_cola.to_dense())
'''
tensor([[0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.2643, 0.0000, 0.0000],
        [0.0000, 0.8338, 0.2440, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.3951]], dtype=torch.float64)
'''
print(coo_tensor_torch.to_dense())
'''
tensor([[0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.2643, 0.0000, 0.0000],
        [0.0000, 0.2440, 0.8338, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.3951]], dtype=torch.float64)
'''

Expected Behavior

The two tensors should be equal but they are not.

System information

Please complete the following information:

• CoLA version 0.0.6
• PyTorch version 2.4.0
• macOS 14.1