Repository: pytorch
Issue: sparse_csr_tensor matmul wrong output in bfloat16
Labels: module: sparse,triaged

### 🐛 Describe the bug

Trying to do a sparse-dense matrix multiplication in with bfloat16 dtype produces incorrect output.

More specifically, I would like to use torch.sparse_csr_tensor to implement conditional matrix multiplication (I have an index selecting which matrix to use, in the sample below I'm simulating this by multiplying each slice of the batch with a different matrix. sparse_bsc_tensor would be more optimal for this, but unfortunately, matmul is not supported with it currently).

The code below tests the output of the operation for different data types and prints the max relative and absolute error:
```python
import torch

n_blocks = 16
n_channels = 128
block_size = 8
bs = 32

device = torch.device("cuda")

# Generate test data
testvec_g = torch.randn(bs, n_channels)
mat_g = torch.randn(n_blocks, n_channels, block_size)

# calculate reference output
olist = []
sample_per_block = bs // n_blocks

for b in range(n_blocks):
    olist.append(testvec_g.float()[b*sample_per_block: (b+1)*sample_per_block] @ mat_g[b].float())

reference = torch.cat(olist, dim=0)

# Do the test
def test(dtype):
    mat = mat_g.to(dtype).to(device)
    testvec = testvec_g.to(dtype).to(device)

    sel = torch.arange(n_blocks)[..., None].expand(-1, bs // n_blocks).flatten().to(device)

    out = torch.sparse_csr_tensor(
        crow_indices=torch.arange(0, testvec.nelement() + 1, testvec.shape[-1], device=mat.device),
        col_indices=(torch.arange(0, testvec.shape[-1], device=mat.device) + sel[:, None] * n_channels).flatten(),
        values=testvec.flatten(),
        size=(testvec.shape[0], testvec.shape[1] * n_blocks)
    )

    scores = out @ mat.flatten(0,1)

    ref = reference.type(dtype).to(device)
    print("  Max forward abs diff: ", (scores - ref).abs().max())
    print("  Max forward rel diff: ", ((scores - ref) / ref).abs().max())


print("float32")
test(torch.float32)

print("float16")
test(torch.float16)

print("bfloat16")
test(torch.bfloat16)
```

It produces the following output (which varies due to the randomly generated inputs):
```
float32
/home/robert/test/min_bfloat16_fail.py:30: UserWarning: Sparse CSR tensor support is in beta state. If you miss a functionality in the sparse tensor support, please submit a feature request to https://github.com/pytorch/pytorch/issues. (Triggered internally at ../aten/src/ATen/SparseCsrTensorImpl.cpp:54.)
  out = torch.sparse_csr_tensor(
  Max forward abs diff:  tensor(7.6294e-06, device='cuda:0')
  Max forward rel diff:  tensor(1.8543e-05, device='cuda:0')
float16
  Max forward abs diff:  tensor(0.0312, device='cuda:0', dtype=torch.float16)
  Max forward rel diff:  tensor(0.1309, device='cuda:0', dtype=torch.float16)
bfloat16
  Max forward abs diff:  tensor(0.2500, device='cuda:0', dtype=torch.bfloat16)
  Max forward rel diff:  tensor(2.0156, device='cuda:0', dtype=torch.bfloat16)
```

Bfloat16 has max 200% error, and just by looking at the matrices by hand, it is clear that the output has nothing like the correct one.

Thanks in advance for any suggestions!

### Versions

```
PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7

OS: Manjaro Linux (x86_64)
GCC version: (GCC) 12.2.0
Clang version: 14.0.6
CMake version: version 3.25.0

Python version: 3.10 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 525.60.11
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.6.0
/usr/lib/libcudnn_adv_infer.so.8.6.0
/usr/lib/libcudnn_adv_train.so.8.6.0
/usr/lib/libcudnn_cnn_infer.so.8.6.0
/usr/lib/libcudnn_cnn_train.so.8.6.0
/usr/lib/libcudnn_ops_infer.so.8.6.0
/usr/lib/libcudnn_ops_train.so.8.6.0

Versions of relevant libraries:
[pip3] kmeans-pytorch==0.3
[pip3] numpy==1.23.5
[pip3] pytorch-lightning==1.9.0
[pip3] pytorch-triton==2.0.0+0d7e753227
[pip3] torch==2.0.0
[pip3] torch-dct==0.1.6
[pip3] torchaudio==2.0.0.dev20230131+cu118
[pip3] torchmetrics==0.11.0
[pip3] torchvision==0.15.0.dev20230131+cu118
[conda] Could not collect
```     

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer