# Owner(s): ["module: sparse"]

import torch
import random
import io
import itertools
import unittest
import functools
from contextlib import redirect_stderr
from torch.testing import make_tensor, FileCheck
from torch.testing._internal.common_cuda import SM53OrLater, SM80OrLater, TEST_CUSPARSE_GENERIC
from torch.testing._internal.common_utils import \
    (TEST_WITH_TORCHINDUCTOR, TEST_WITH_ROCM, TEST_CUDA_CUDSS, TEST_SCIPY, TEST_NUMPY, TEST_MKL, IS_WINDOWS, TestCase,
     run_tests, load_tests, coalescedonoff, parametrize, subtest, skipIfTorchDynamo, skipIfRocm, IS_FBCODE, IS_REMOTE_GPU,
     suppress_warnings)
from torch.testing._internal.common_device_type import \
    (ops, instantiate_device_type_tests, dtypes, OpDTypes, dtypesIfCUDA, onlyCPU, onlyCUDA, skipCUDAIfNoSparseGeneric,
     precisionOverride, skipMeta, skipCUDAIf, skipCPUIfNoMklSparse, skipCUDAIfRocmVersionLessThan,
     largeTensorTest)
from torch.testing._internal.common_methods_invocations import \
    (op_db, sparse_csr_unary_ufuncs, ReductionOpInfo)
from torch.testing._internal.common_cuda import _get_torch_cuda_version, TEST_CUDA
from torch.testing._internal.common_dtype import (
    floating_types, all_types_and_complex_and, floating_and_complex_types, floating_types_and,
    all_types_and_complex, floating_and_complex_types_and)
from torch.testing._internal.opinfo.definitions.linalg import sample_inputs_linalg_solve
from torch.testing._internal.opinfo.definitions.sparse import validate_sample_input_sparse
from test_sparse import CUSPARSE_SPMM_COMPLEX128_SUPPORTED, HIPSPARSE_SPMM_COMPLEX128_SUPPORTED
import operator

if TEST_SCIPY:
    import scipy.sparse as sp

if TEST_NUMPY:
    import numpy as np
# load_tests from torch.testing._internal.common_utils is used to automatically filter tests for
# sharding on sandcastle. This line silences flake warnings
load_tests = load_tests

no_mkl_sparse = IS_WINDOWS or not TEST_MKL

def _check_cusparse_triangular_solve_available():
    version = _get_torch_cuda_version()
    # cusparseSpSM was added in 11.3.1 but we don't have access to patch version
    min_supported_version = (11, 4)
    return version >= min_supported_version

def _check_cusparse_spgemm_available():
    # cusparseSpGEMM was added in 11.0
    return not TEST_WITH_ROCM

def _check_cusparse_sddmm_available():
    if TEST_WITH_ROCM:
        return True
    version = _get_torch_cuda_version()
    # cusparseSDDMM was added in 11.2.1 but we don't have access to patch version
    min_supported_version = (11, 3)
    return version >= min_supported_version

_sparse_csr_ops = list(filter(lambda op: op.supports_sparse_csr, op_db))
_sparse_compressed_ops = list(filter(lambda op: (op.supports_sparse_csr or op.supports_sparse_csc
                                                 or op.supports_sparse_bsr or op.supports_sparse_bsc), op_db))
binary_functions_with_dense_output = ['mm', 'mv', ]
binary_ops_with_dense_output = list(filter(lambda op: op.name in binary_functions_with_dense_output, op_db))

UNARY_EWISE_CSR_ALLOW_AUTOGRAD = [
    'abs',
    'conj_physical',
    'deg2rad',
    'neg',
    'positive',
    'frac',
    'nn.functional.relu',
    'log1p',
    'rad2deg'
]

# This should be just an import from test_linalg instead of code duplication
# but https://github.com/pytorch/pytorch/pull/63511#discussion_r733989701
def _test_addmm_addmv(
    test_case,
    f,
    t,
    m,
    v,
    *,
    alpha=None,
    beta=None,
    transpose_out=False,
    layout=torch.strided,
    mode=None
):
    """
    Unified test for checking `f(t, m, v, alpha=alpha, beta=beta)` computation,
    where f is `torch.addmv` or `torch.addmm`.
    `transpose_out` controls whether the out argument is in column-major order.
    `layout` controls whether `m` is converted to specified layout or not.
    Custom behaviour is implemented only for torch.sparse_csr layout.
    """
    dtype = t.dtype
    numpy_dtype = dtype
    if dtype in {torch.bfloat16}:
        numpy_dtype = torch.float
    if dtype.is_complex:
        alpha = 0.9 + 0.3j if alpha is None else alpha
        beta = 0.5 + 0.6j if beta is None else beta
    else:
        alpha = 1.2 if alpha is None else alpha
        beta = 0.8 if beta is None else beta

    def convert_layout(mat):
        if layout == torch.sparse_csr:
            return mat.to_sparse_csr()
        elif layout == torch.sparse_csc:
            return mat.to_sparse_csc()
        else:
            assert mat.layout == layout
            return mat

    if mode == "all_sparse":
        res1 = f(*map(convert_layout, (t, m, v)), alpha=alpha, beta=beta)
        test_case.assertEqual(res1.layout, layout)
        res1 = res1.to_dense()
    elif mode == "dense_result":
        res1 = f(t, convert_layout(m), convert_layout(v), alpha=alpha, beta=beta)
    else:
        res1 = f(t, convert_layout(m), v, alpha=alpha, beta=beta)
    res2 = torch.full_like(res1, float('nan'))
    if transpose_out:
        res2 = res2.t().clone(memory_format=torch.contiguous_format).t()
    f(t, convert_layout(m), v, alpha=alpha, beta=beta, out=res2)
    res3 = alpha * (m.to(numpy_dtype).cpu().numpy() @ v.to(numpy_dtype).cpu().numpy())
    if beta != 0:
        res3 += (beta * t).to(numpy_dtype).cpu().numpy()
    res3 = torch.from_numpy(res3).to(dtype)
    test_case.assertEqual(res1, res2)
    test_case.assertEqual(res1, res3)


class TestSparseCSRSampler(TestCase):

    def test_make_crow_indices(self):
        # Here we test the correctness of the crow_indices algorithm
        # and testing it on CPU and with int32 dtype will be
        # sufficient.
        device = torch.device('cpu')
        index_dtype = torch.int32
        for n_rows in range(1, 10):
            for n_cols in range(1, 10):
                for nnz in range(0, n_rows * n_cols + 1):
                    crow_indices = self._make_crow_indices(
                        n_rows, n_cols, nnz,
                        device=device, dtype=index_dtype)
                    self.assertEqual(len(crow_indices), n_rows + 1)
                    counts = crow_indices[1:] - crow_indices[:-1]
                    self.assertEqual(counts.sum(), nnz)
                    self.assertGreaterEqual(counts.min(), 0)
                    self.assertLessEqual(counts.max(), n_cols)


def all_sparse_compressed_layouts(test_name='layout'):
    return parametrize(test_name, [
        subtest(torch.sparse_csr, name='SparseCSR'),
        subtest(torch.sparse_csc, name='SparseCSC'),
        subtest(torch.sparse_bsr, name='SparseBSR'),
        subtest(torch.sparse_bsc, name='SparseBSC')])


def sparse_compressed_nonblock_layouts(test_name='layout'):
    return parametrize(test_name, [
        subtest(torch.sparse_csr, name='SparseCSR'),
        subtest(torch.sparse_csc, name='SparseCSC')])


sparse_compressed_indices_methods = {
    torch.sparse_csr: (torch.Tensor.crow_indices, torch.Tensor.col_indices),
    torch.sparse_csc: (torch.Tensor.ccol_indices, torch.Tensor.row_indices),
    torch.sparse_bsr: (torch.Tensor.crow_indices, torch.Tensor.col_indices),
    torch.sparse_bsc: (torch.Tensor.ccol_indices, torch.Tensor.row_indices),
}


def batched_nonbatched(test_name='batched'):
    return parametrize(test_name, [
        subtest(True, name="Batched"),
        subtest(False, name="NonBatched")
    ])


def hybrid_nonhybrid(test_name='hybrid'):
    return parametrize(test_name, [
        subtest(True, name="Hybrid"),
        subtest(False, name="NonHybrid")
    ])


class TestSparseCompressed(TestCase):
    """Testing sparse compressed (CSR, CSC, BSR, BSC) tensor generic features.
    """

    def genTensor(self, size, nnz, *, layout, device=None, dtype=torch.float, index_dtype=torch.int64):
        if device is None:
            device = self.device_type
        return self.genSparseCompressedTensor(size, nnz, device=device, dtype=dtype, index_dtype=index_dtype, layout=layout)

    @all_sparse_compressed_layouts()
    @onlyCPU
    def test_layout(self, layout):
        self.assertIn(str(layout), {'torch.sparse_csr', 'torch.sparse_csc', 'torch.sparse_bsr', 'torch.sparse_bsc'})
        self.assertEqual(type(layout), torch.layout)

    @parametrize('shape_and_device_inference', [subtest(False, name='_'), subtest(True, name='shape_and_device_inference')])
    @parametrize('use_factory_function', [subtest(False, name='_'), subtest(True, name='factory')])
    @parametrize('input_kind', [subtest('tensor', name='from_tensor'), subtest('list', name='from_list')])
    @all_sparse_compressed_layouts()
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_sparse_compressed_constructor(self, layout, device, dtype,
                                           use_factory_function, shape_and_device_inference, input_kind):
        if input_kind == 'list' and shape_and_device_inference:
            if torch.device(device).type == 'cuda':
                # list inputs to factory/constructor function without
                # specifying device will result a sparse compressed tensor
                # on CPU. So, skip testing against cuda device as unused.
                self.skipTest("nothing to test")
            if dtype not in {torch.float32, torch.complex64, torch.int64, torch.bool}:
                self.skipTest("dtype not supported with list values")

        expected_devices = [torch.device(device)]
        if TEST_CUDA and torch.device(device).type == 'cuda' and torch.cuda.device_count() >= 2 and not shape_and_device_inference:
            expected_devices.append(torch.device('cuda:1'))

        factory_function = {
            torch.sparse_csr: torch.sparse_csr_tensor,
            torch.sparse_csc: torch.sparse_csc_tensor,
            torch.sparse_bsr: torch.sparse_bsr_tensor,
            torch.sparse_bsc: torch.sparse_bsc_tensor,
        }[layout]
        compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[layout]
        if input_kind == 'list':
            index_dtypes = [torch.int64]
        else:
            index_dtypes = [torch.int32, torch.int64]
        if dtype.is_floating_point or dtype.is_complex:
            requires_grad_lst = [False, True]
        else:
            requires_grad_lst = [False]
        for index_dtype in index_dtypes:
            for expected_device in expected_devices:
                for (compressed_indices, plain_indices, values), kwargs in self.generate_simple_inputs(
                        layout, device=expected_device, dtype=dtype, index_dtype=index_dtype,
                        # skip zero-sized tensors for list inputs:
                        enable_zero_sized=input_kind != 'list',
                        output_tensor=False):
                    size = kwargs['size']
                    if shape_and_device_inference and 0 in size:
                        # skip shape inference for zero-sized tensor
                        # inputs because (i) the shape determined from
                        # an empty list is ambiguous, and (ii) the
                        # size of the plain dimension defined as
                        # max(plain_indices) is undefined if
                        # plain_indices has no values
                        continue
                    compressed_indices_expect = compressed_indices
                    plain_indices_expect = plain_indices
                    values_expect = values

                    if input_kind == 'list':
                        compressed_indices = compressed_indices.tolist()
                        plain_indices = plain_indices.tolist()
                        values = values.tolist()

                    for requires_grad in requires_grad_lst:
                        if use_factory_function:
                            if shape_and_device_inference:
                                sparse = factory_function(
                                    compressed_indices, plain_indices, values, requires_grad=requires_grad)
                            else:
                                sparse = factory_function(
                                    compressed_indices, plain_indices, values, size,
                                    dtype=dtype, device=expected_device, requires_grad=requires_grad)
                        else:
                            if shape_and_device_inference:
                                sparse = torch.sparse_compressed_tensor(
                                    compressed_indices, plain_indices, values,
                                    layout=layout, requires_grad=requires_grad)
                            else:
                                sparse = torch.sparse_compressed_tensor(
                                    compressed_indices, plain_indices, values, size,
                                    dtype=dtype, layout=layout, device=expected_device, requires_grad=requires_grad)

                        self.assertEqual(layout, sparse.layout)
                        self.assertEqual(size, sparse.shape)
                        self.assertEqual(compressed_indices_expect, compressed_indices_mth(sparse))
                        self.assertEqual(plain_indices_expect, plain_indices_mth(sparse))
                        self.assertEqual(values_expect, sparse.values())
                        self.assertEqual(sparse.device, sparse.values().device)
                        self.assertEqual(sparse.device, expected_device)
                        self.assertEqual(sparse.values().requires_grad, requires_grad)
                        self.assertEqual(sparse.requires_grad, requires_grad)
                        self.assertFalse(compressed_indices_mth(sparse).requires_grad)
                        self.assertFalse(plain_indices_mth(sparse).requires_grad)

    @skipMeta
    @sparse_compressed_nonblock_layouts()
    @dtypes(*all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half))
    def test_empty(self, layout, device, dtype):
        ns = [5, 2, 0]
        batch_shapes = [(), (2,), (2, 3)]
        compressed_dim = {
            torch.sparse_csr: -2,
            torch.sparse_csc: -1,
        }[layout]
        compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[layout]
        for m, n, b in itertools.product(ns, ns, batch_shapes):
            shape = (*b, m, n)
            with torch.sparse.check_sparse_tensor_invariants(enable=False):
                # torch.empty may return invalid sparse compressed tensors
                result = torch.empty(shape, dtype=dtype, device=device, layout=layout)
            self.assertEqual(result.shape, shape)
            self.assertEqual(result.dtype, dtype)
            self.assertEqual(result.device, torch.device(device))
            self.assertEqual(result.layout, layout)
            self.assertEqual(compressed_indices_mth(result).shape, (*b, shape[compressed_dim] + 1,))
            self.assertEqual(plain_indices_mth(result).shape, (*b, 0,))
            self.assertEqual(result.values().shape, (*b, 0,))
            self.assertEqual(result._nnz(), 0)
            self.assertEqual(compressed_indices_mth(result).device, torch.device(device))
            self.assertEqual(plain_indices_mth(result).device, torch.device(device))
            self.assertEqual(result.values().device, torch.device(device))
            self.assertEqual(compressed_indices_mth(result).dtype, torch.int64)
            self.assertEqual(plain_indices_mth(result).dtype, torch.int64)
            self.assertEqual(result.values().dtype, dtype)

    @skipMeta
    @sparse_compressed_nonblock_layouts()
    @dtypes(*all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16))
    def test_empty_errors(self, layout, device, dtype):
        with self.assertRaisesRegex(RuntimeError,
                                    "torch.empty: Only batched sparse compressed \\(non-block\\) tensors are supported"
                                    ", but got size"):
            torch.empty((5,), dtype=dtype, device=device, layout=layout)

    @skipMeta
    @all_sparse_compressed_layouts()
    @dtypes(*all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half))
    def test_sparse_compressed_tensor_with_dims(self, layout, device, dtype):

        def get_sparse_compressed_tensor_properties(s):
            if layout in {torch.sparse_csr, torch.sparse_bsr}:
                compressed_indices, plain_indices = s.crow_indices(), s.col_indices()
            else:
                compressed_indices, plain_indices = s.ccol_indices(), s.row_indices()
            values = s.values()
            return dict(shape=s.shape, dtype=s.dtype, device=s.device, nnz=s._nnz(), layout=s.layout,
                        compressed_indices_shape=compressed_indices.shape,
                        compressed_indices_dtype=compressed_indices.dtype,
                        compressed_indices_device=compressed_indices.device,
                        plain_indices_shape=plain_indices.shape,
                        plain_indices_dtype=plain_indices.dtype,
                        plain_indices_device=plain_indices.device,
                        values_shape=values.shape,
                        values_dtype=values.dtype,
                        values_device=values.device)

        for index_dtype in [torch.int32, torch.int64]:
            for t in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype):
                dense_dim = t.dense_dim()
                sparse_dim = t.sparse_dim()
                batch_dim = t.ndim - sparse_dim - dense_dim
                nnz = t.values().shape[batch_dim]
                if layout in {torch.sparse_bsr, torch.sparse_bsc}:
                    blocksize = t.values().shape[batch_dim + 1: batch_dim + 1 + sparse_dim]
                else:
                    blocksize = ()

                e = torch.ops.aten._sparse_compressed_tensor_with_dims(nnz, dense_dim, t.shape, blocksize, index_dtype,
                                                                       dtype=dtype, layout=layout, device=device)

                e_prop, t_prop = get_sparse_compressed_tensor_properties(e), get_sparse_compressed_tensor_properties(t)
                for k, v in e_prop.items():
                    self.assertEqual(v, t_prop[k], lambda msg: f'{msg} when comparing {k}, expected {t_prop[k]}, got {v}')

    @skipMeta
    @all_sparse_compressed_layouts()
    @dtypes(*all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16))
    def test_clone(self, layout, device, dtype):
        for sparse in self.generate_simple_inputs(
                layout, device=device, dtype=dtype, index_dtype=torch.int32):
            cloned_sparse = sparse.clone()
            self.assertEqual(sparse, cloned_sparse)

    @all_sparse_compressed_layouts()
    def test_print(self, layout, device):
        compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[layout]
        printed = []
        for enable_hybrid in [False, True]:
            # using local patterns for test_print stability
            patterns = [
                # 2 x 3 batch of 3 x 2 tensors, trivial blocksize, non-hybrid/hybrid:
                ([[[[1, 2, 0],
                    [1, 0, 3]],
                   [[1, 2, 3],
                    [1, 0, 0]],
                   [[1, 0, 0],
                    [1, 2, 3]]],
                  [[[0, 2, 0],
                    [1, 2, 3]],
                   [[1, 0, 3],
                    [1, 2, 0]],
                   [[1, 2, 3],
                    [0, 2, 0]]]], [(2, 1)], [(), (4,)] if enable_hybrid else [()]),
                # tensor with non-trivial blocksize, non-hybrid/hybrid:
                ([[0, 1, 0, 2, 0, 2],
                  [0, 1, 0, 0, 2, 0],
                  [3, 3, 3, 0, 0, 0],
                  [0, 0, 0, 0, 0, 0],
                  [0, 5, 0, 6, 6, 6],
                  [5, 0, 5, 6, 6, 6],
                  [0, 0, 0, 0, 8, 8],
                  [7, 7, 7, 0, 8, 8]], [(2, 3)], [(), (4, 2)] if enable_hybrid else [()]),
            ]
            for index_dtype in [torch.int32, torch.int64]:
                for dtype in [torch.float32, torch.float64]:
                    for (compressed_indices, plain_indices, values), kwargs in self.generate_simple_inputs(
                            layout, device=device, dtype=dtype, index_dtype=index_dtype, enable_hybrid=enable_hybrid,
                            enable_non_contiguous_indices=False, enable_non_contiguous_values=False,
                            enable_zero_sized=False, output_tensor=False, patterns=patterns):
                        size = tuple(kwargs['size'])
                        block_ndim = 2 if layout in {torch.sparse_bsr, torch.sparse_bsc} else 0
                        base_ndim = 2
                        batch_ndim = compressed_indices.dim() - 1
                        dense_ndim = values.dim() - batch_ndim - block_ndim - 1
                        if enable_hybrid and dense_ndim == 0:
                            # non-hybrid cases are covered by the enable_hybrid==False loop
                            continue
                        batchsize = size[:batch_ndim]
                        basesize = size[batch_ndim:batch_ndim + base_ndim]
                        densesize = size[batch_ndim + base_ndim:]
                        assert len(densesize) == dense_ndim
                        printed.append(f"########## {dtype}/{index_dtype}/size={batchsize}+{basesize}+{densesize} ##########")
                        x = torch.sparse_compressed_tensor(compressed_indices,
                                                           plain_indices,
                                                           values, size, dtype=dtype, layout=layout, device=device)
                        printed.append("# sparse tensor")
                        printed.append(str(x))
                        printed.append(f"# _{compressed_indices_mth.__name__}")
                        printed.append(str(compressed_indices_mth(x)))
                        printed.append(f"# _{plain_indices_mth.__name__}")
                        printed.append(str(plain_indices_mth(x)))
                        printed.append("# _values")
                        printed.append(str(x.values()))
                        printed.append('')
                    printed.append('')
        orig_maxDiff = self.maxDiff
        self.maxDiff = None
        try:
            self.assertExpected('\n'.join(printed))
            self.maxDiff = orig_maxDiff
        except Exception:
            self.maxDiff = orig_maxDiff
            raise

    @skipMeta
    @all_sparse_compressed_layouts()
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_copy(self, layout, device, dtype):

        def run_test(shape, blocksize, nnz, index_type):
            a = self.genSparseCompressedTensor(shape, nnz, dtype=dtype, layout=layout, device=device,
                                               index_dtype=index_dtype, blocksize=blocksize)
            b = self.genSparseCompressedTensor(shape, nnz, dtype=dtype, layout=layout, device=device,
                                               index_dtype=index_dtype, blocksize=blocksize)

            a.copy_(b)

            self.assertEqual(a, b)

        ns = [(9, 3), (2, 1), (0, 0)]  # (number of dimensions, the corresponding block size)
        batch_shapes = [(), (2,), (2, 3)]
        for ((m, bm), (n, bn), b), index_dtype in zip(itertools.product(ns, ns, batch_shapes), [torch.int32, torch.int64]):
            blocksize = (bm, bn) if layout in {torch.sparse_bsr, torch.sparse_bsc} else ()
            run_test((*b, m, n), blocksize, 0, index_dtype)
            run_test((*b, m, n), blocksize, m * n, index_dtype)

    @skipMeta
    @all_sparse_compressed_layouts()
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_copy_errors(self, layout, device, dtype):
        blocksize = (2, 3) if layout in {torch.sparse_bsr, torch.sparse_bsc} else ()
        nnz = 6 if layout in {torch.sparse_bsr, torch.sparse_bsc} else 1
        shape1 = (2 * 6, 3 * 6) if layout in {torch.sparse_bsr, torch.sparse_bsc} else (2, 3)
        for index_dtype in [torch.int32, torch.int64]:
            a = self.genSparseCompressedTensor(shape1, 0, dtype=dtype, layout=layout, device=device,
                                               index_dtype=index_dtype, blocksize=blocksize)

            with self.assertRaisesRegex(RuntimeError,
                                        "copy of sparse compressed tensors having different layouts is not supported."):
                a.copy_(torch.empty(a.shape, dtype=dtype, device=device))

            b = self.genSparseCompressedTensor(shape1, nnz, dtype=dtype, layout=layout, device=device,
                                               index_dtype=index_dtype, blocksize=blocksize)
            assert a._nnz() != b._nnz(), (a._nnz(), b._nnz())
            with self.assertRaisesRegex(RuntimeError,
                                        "only sparse compressed tensors with the same number of specified elements are supported."):
                a.copy_(b)

            shape2 = tuple(reversed(shape1))
            c = self.genSparseCompressedTensor(shape2, nnz, dtype=dtype, layout=layout, device=device,
                                               index_dtype=index_dtype, blocksize=blocksize)
            with self.assertRaisesRegex(
                    RuntimeError,
                    "expected shapes of self and src to match along dimension"):
                b.copy_(c)

            if blocksize:
                blocksize1 = tuple(reversed(blocksize))
                d = self.genSparseCompressedTensor(shape1, nnz, dtype=dtype, layout=layout, device=device,
                                                   index_dtype=index_dtype, blocksize=blocksize1)
                with self.assertRaisesRegex(RuntimeError,
                                            "copy of sparse compressed tensors having different block sizes is not supported"):
                    b.copy_(d)

    def _smallest_divisor(self, n):
        for i in range(2, int(n ** 0.5) + 1):
            if n % i == 0:
                return i
        return n

    @skipIfTorchDynamo("Not a TorchDynamo suitable test")
    @all_sparse_compressed_layouts()
    @ops(_sparse_compressed_ops)
    @precisionOverride({torch.bfloat16: 1e-2, torch.float16: 1e-2})
    def test_consistency(self, layout, device, dtype, op):
        """Checks that the op on a strided and on a sparse tensors will
        produce the same results.
        """
        if not op.supports_sparse_layout(layout):
            self.skipTest(f"{op.name} does not support input with {layout} layout")

        # FIXME: remove in followup once integer support is landed for segment_reduce
        if (layout == torch.sparse_csr and not dtype.is_floating_point
                and op.name in ('masked.mean', 'masked.amax', 'masked.amin')):
            self.skipTest(f"{op.name} does not support input with {layout} layout and {dtype} dtype")

        require_mask = isinstance(op, ReductionOpInfo) and 'masked.' in op.name

        samples = []
        for sample in op.sample_inputs(device, dtype):
            if sample.input.ndim < 2:
                continue
            dense_dim = sample.input.ndim - 2
            blocksize = (tuple(map(self._smallest_divisor, sample.input.shape[:2]))
                         if layout in {torch.sparse_bsr, torch.sparse_bsc} else None)

            def _to_sparse(x):
                if isinstance(x, torch.Tensor):
                    if blocksize is None:
                        if x.ndim != sample.input.ndim:
                            return x
                    elif x.ndim != sample.input.ndim + 2 or x.shape[-3] % blocksize[0] or x.shape[-2] % blocksize[1]:
                        return x
                    return x.clone().to_sparse(layout=layout, blocksize=blocksize, dense_dim=dense_dim)
                return x

            sparse_sample = sample.transform(_to_sparse)
            # Some strided samples (with inf, nan elements) appear to share
            # storage, so we must clone:
            sample = sample.transform(lambda x: (x.clone() if isinstance(x, torch.Tensor) else x))

            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.")

        # Re-define atol and rtol for operations that result values
        # are random (and hence, non-comparable) be we still want to
        # check the shape, dtype, etc attributes of the results:
        atol = rtol = None
        if op.name == 'randn_like':
            atol = 1e300
            rtol = 1

        for sample, sparse_sample in samples:
            expected = op(sample.input, *sample.args, **sample.kwargs)
            assert torch.is_tensor(expected)
            output = op(sparse_sample.input, *sparse_sample.args, **sparse_sample.kwargs)
            assert torch.is_tensor(output)
            strided_output = output.to_dense()
            if require_mask and sample.kwargs.get('mask') is not None:
                output_mask = torch.masked._output_mask(op.op, sample.input, *sample.args, **sample.kwargs)
                expected.masked_fill_(~output_mask, 0)
            self.assertEqual(strided_output, expected, atol=atol, rtol=rtol)

    @skipMeta
    @all_sparse_compressed_layouts()
    @all_sparse_compressed_layouts('layout2')
    @dtypes(*all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16))
    def test_empty_like(self, layout, layout2, device, dtype):
        for sparse in self.generate_simple_inputs(layout):
            if layout == layout2:
                result = torch.empty_like(sparse, layout=layout2)
                compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[result.layout]
                torch._validate_sparse_compressed_tensor_args(compressed_indices_mth(result),
                                                              plain_indices_mth(result),
                                                              result.values(),
                                                              result.shape,
                                                              result.layout)
                self.assertEqual(sparse.shape, result.shape)
            else:
                self.assertRaisesRegex(
                    RuntimeError,
                    "empty_like with different sparse layout is not supported",
                    lambda: torch.empty_like(sparse, layout=layout2)
                )

    @skipMeta
    @all_sparse_compressed_layouts()
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_validate(self, layout, device, dtype):
        def make_zero_batched(t):
            return torch.empty(*((0,) + t.shape), dtype=t.dtype, device=t.device)

        for index_dtype in [torch.int32, torch.int64]:
            for (compressed_indices, plain_indices, values), kwargs in self.generate_simple_inputs(
                    layout, device=device, dtype=dtype, index_dtype=index_dtype, output_tensor=False):
                size = kwargs['size']
                torch._validate_sparse_compressed_tensor_args(compressed_indices, plain_indices, values, size, layout)

                # check empty batch
                torch._validate_sparse_compressed_tensor_args(
                    *(make_zero_batched(t) for t in (compressed_indices, plain_indices, values)),
                    (0,) + size,
                    layout
                )

            compressed_indices = torch.tensor([0, 0], dtype=index_dtype)
            plain_indices = torch.tensor([], dtype=index_dtype)
            torch._validate_compressed_sparse_indices(layout in {torch.sparse_csr, torch.sparse_bsr},
                                                      compressed_indices, plain_indices, 1, 1, 0)

    def _generate_invalid_input(self, layout, device):
        from functools import partial

        def shape(shape, basedim=0):
            blocksize = (1, 1)
            if layout is torch.sparse_csc:
                shape = shape[:basedim] + (shape[basedim + 1], shape[basedim]) + shape[basedim + 2:]
            elif layout is torch.sparse_bsc:
                shape = shape[:basedim] + (shape[basedim + 1] * blocksize[1], shape[basedim] * blocksize[0]) + shape[basedim + 2:]
            elif layout is torch.sparse_bsr:
                shape = shape[:basedim] + (shape[basedim] * blocksize[0], shape[basedim + 1] * blocksize[1]) + shape[basedim + 2:]
            return shape

        def values(lst, device=device):
            if layout in {torch.sparse_bsr, torch.sparse_bsc}:
                lst = [[[item]] for item in lst]
            return torch.tensor(lst, device=device)

        tensor = partial(torch.tensor, device=device)
        values = partial(values, device=device)

        yield ('incontiguous compressed_indices',
               tensor([0, -1, 2, -1, 4, -1])[::2],
               tensor([0, 1, 0, 2]),
               values([1, 2, 3, 4]),
               shape((2, 3)),
               'expected compressed_indices to be a contiguous tensor per batch')

        yield ('incontiguous plain_indices',
               tensor([0, 2, 4]),
               tensor([0, -1, 1, -1, 0, -1, 2, -1])[::2],
               values([1, 2, 3, 4]),
               shape((2, 3)),
               'expected plain_indices to be a contiguous tensor per batch')

        yield ('0-D compressed_indices',
               tensor(0),
               tensor([0, 1, 0, 2]),
               values([1, 2, 3, 4]),
               shape((2, 3)),
               'compressed_indices must have dimensionality >= 1 but got 0')

        yield ('compressed/plain_indices mismatch of dimensionalities',
               tensor([[0, 2, 4]]),
               tensor([0, 1, 0, 2]),
               values([1, 2, 3, 4]),
               shape((2, 3)),
               'compressed_indices and plain_indices dimensionalities must be equal but got 2 and 1, respectively')

        if layout in {torch.sparse_csr, torch.sparse_csc}:
            yield ('indices and values mismatch of dimensionalities',
                   tensor([[0, 2, 4]]),
                   tensor([[0, 1, 0, 2]]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'values must have dimensionality > sum of batch and block dimensionalities \(=1 \+ 0\) but got 1')
        else:
            yield ('indices and values mismatch of dimensionalities',
                   tensor([[0, 2, 4]]),
                   tensor([[0, 1, 0, 2]]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'values must have dimensionality > sum of batch and block dimensionalities \(=1 \+ 2\) but got 3')

        yield ('invalid size',
               tensor([0, 2, 4]),
               tensor([0, 1, 0, 2]),
               values([1, 2, 3, 4]),
               (2,),
               r'tensor dimensionality must be sum of batch, base, and dense dimensionalities \(=0 \+ 2 \+ 0\) but got 1')

        yield ('invalid batchsize',
               tensor([[0, 2, 4]]),
               tensor([[0, 1, 0, 2]]),
               values([[1, 2, 3, 4]]),
               shape((2, 2, 3), 1),
               r'all batch dimensions of compressed_indices \(=\[1\]\), plain_indices \(=\[1\]\), '
               r'and values \(=\[1\]\) must be equal to tensor batch dimensions \(=\[2\]\)')

        if layout is torch.sparse_bsr:
            yield ('invalid blocksize',
                   tensor([0, 2, 4]),
                   tensor([0, 1, 0, 2]),
                   tensor([[[1, 11]], [[2, 22]], [[3, 33]], [[4, 33]]]),
                   shape((2, 3)),
                   r'tensor shape\[1\] \(=3\) must be divisible with blocksize\[1\] \(=2\) as defined by values shape')

        if layout is torch.sparse_bsc:
            yield ('invalid blocksize',
                   tensor([0, 2, 4]),
                   tensor([0, 1, 0, 2]),
                   tensor([[[1, 11]], [[2, 22]], [[3, 33]], [[4, 33]]]),
                   shape((3, 2)),
                   r'tensor shape\[1\] \(=3\) must be divisible with blocksize\[1\] \(=2\) as defined by values shape')

        yield ('invalid compressed_indices shape',
               tensor([0, 2, 3, 4]),
               tensor([0, 1, 0, 2]),
               values([1, 2, 3, 4]),
               shape((2, 3)),
               r'compressed_indices.shape\[-1\] must be equal to the number of compressed_indices_names \+ 1 \(=3\), but got 4')

        yield ('invalid compressed_indices shape',
               tensor([0, 2, 4]),
               tensor([0, 1, 0, 1, 2]),
               values([1, 2, 3, 4]),
               shape((2, 3)),
               r'plain_indices.shape\[-1\] must be equal to nnz \(=4\) as defined by values.shape\[0\], but got 5')

        yield ('compressed/plain_indices mismatch of dtype',
               tensor([0, 2, 4], dtype=torch.int32),
               tensor([0, 1, 0, 2], dtype=torch.int64),
               values([1, 2, 3, 4]),
               shape((2, 3)),
               r'compressed_indices and plain_indices must have the same dtype, bot got Int and Long, respectively')

        yield ('invalid compressed/plain_indices dtype',
               tensor([0, 2, 4], dtype=torch.int16),
               tensor([0, 1, 0, 2], dtype=torch.int16),
               values([1, 2, 3, 4]),
               shape((2, 3)),
               r'compressed_indices and plain_indices dtype must be Int or Long, but got Short')

        # CUDA kernel asserts are not recoverable, so we skip these for now
        if torch.device(device).type == 'cpu':
            yield ('invalid compressed_indices[0]',
                   tensor([1, 2, 4]),
                   tensor([0, 1, 0, 2]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'`compressed_indices\[..., 0\] == 0` is not satisfied.')

            yield ('invalid compressed_indices[0] when nnz == 0',
                   tensor([1, 0], dtype=torch.int64),
                   tensor([], dtype=torch.int64),
                   values([1])[:0],
                   shape((1, 1)),
                   r'`compressed_indices\[..., 0\] == 0` is not satisfied.')

            yield ('invalid compressed_indices[-1]',
                   tensor([0, 2, 5]),
                   tensor([0, 1, 0, 2]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'`compressed_indices\[..., -1\] == nnz` is not satisfied.')

            yield ('invalid compressed_indices[-1] when nnz == 0',
                   tensor([0, 1], dtype=torch.int64),
                   tensor([], dtype=torch.int64),
                   values([1])[:0],
                   shape((1, 1)),
                   r'`compressed_indices\[..., -1\] == nnz` is not satisfied.')

            yield ('invalid compressed_indices.diff(dim=-1)',
                   tensor([0, 0, 4]),
                   tensor([0, 1, 0, 2]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'0 <= compressed_indices\[..., 1:\] - compressed_indices\[..., :\-1\] <= plain_dim` is not satisfied.')

            yield ('invalid compressed_indices.diff(dim=-1)',
                   tensor([0, 5, 4]),
                   tensor([0, 1, 0, 2]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'0 <= compressed_indices\[..., 1:\] - compressed_indices\[..., :\-1\] <= plain_dim` is not satisfied.')

            yield ('invalid min(plain_indices)',
                   tensor([0, 2, 4]),
                   tensor([0, -1, 0, 3]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'`0 <= plain_indices < plain_dim` is not satisfied.')

            yield ('invalid max(plain_indices)',
                   tensor([0, 2, 4]),
                   tensor([0, 1, 0, 3]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'`0 <= plain_indices < plain_dim` is not satisfied.')

            yield ('non-coalesced',
                   tensor([0, 2, 4]),
                   tensor([1, 0, 0, 2]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'`plain_indices\[..., compressed_indices\[..., i - 1\]:compressed_indices\[..., i\]\] '
                   'for all i = 1, ..., compressed_dim '
                   'are sorted and distinct along the last dimension values` is not satisfied.')

        if TEST_CUDA and torch.device(device).type == 'cpu':
            yield ('indices and values mismatch of device',
                   torch.tensor([0, 2, 4]),
                   torch.tensor([0, 1, 0, 1]),
                   values([1, 2, 3, 4], device='cuda'),
                   shape((2, 3)),
                   r'device of compressed_indices \(=cpu\) must match device of values \(=cuda:0\)')
            yield ('compressed_indices and values mismatch of device',
                   torch.tensor([0, 2, 4], device='cuda'),
                   torch.tensor([0, 1, 0, 1]),
                   values([1, 2, 3, 4]),
                   shape((2, 3)),
                   r'Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!')
            yield ('compressed/plain_indices mismatch of device',
                   torch.tensor([0, 2, 4], device='cuda'),
                   torch.tensor([0, 1, 0, 1]),
                   values([1, 2, 3, 4], device='cuda'),
                   shape((2, 3)),
                   r'Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!')

        if TEST_CUDA and torch.device(device).type == 'cuda' and torch.cuda.device_count() >= 2:
            yield ('indices and values mismatch of device index',
                   torch.tensor([0, 2, 4], device='cuda:0'),
                   torch.tensor([0, 1, 0, 1], device='cuda:0'),
                   values([1, 2, 3, 4], device='cuda:1'),
                   shape((2, 3)),
                   r'device of compressed_indices \(=cuda:0\) must match device of values \(=cuda:1\)')
            yield ('compressed_indices and values mismatch of device index',
                   torch.tensor([0, 2, 4], device='cuda:0'),
                   torch.tensor([0, 1, 0, 1], device='cuda:1'),
                   values([1, 2, 3, 4], device='cuda:0'),
                   shape((2, 3)),
                   r'Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!')

    @skipMeta
    @all_sparse_compressed_layouts()
    @parametrize('target', [subtest('validate_sparse_compressed_tensor_args'),
                            subtest('sparse_compressed_tensor'),
                            subtest('sparse_compressed_tensor_no_size')])
    def test_invalid_input(self, layout, device, target):
        for label, compressed_indices, plain_indices, values, size, errmsg in self._generate_invalid_input(layout, device):
            if layout is torch.sparse_bsr:
                errmsg = errmsg.replace('compressed_indices_name', 'row block').replace('plain_indices_name', 'column block')
            elif layout is torch.sparse_bsc:
                errmsg = errmsg.replace('compressed_indices_name', 'column block').replace('plain_indices_name', 'row block')
            elif layout is torch.sparse_csr:
                errmsg = errmsg.replace('compressed_indices_name', 'row').replace('plain_indices_name', 'column')
            elif layout is torch.sparse_csc:
                errmsg = errmsg.replace('compressed_indices_name', 'column').replace('plain_indices_name', 'row')
            if layout in {torch.sparse_csr, torch.sparse_bsr}:
                errmsg = errmsg.replace('compressed_indices', 'crow_indices') \
                               .replace('plain_indices', 'col_indices') \
                               .replace('plain_dim', 'ncols') \
                               .replace('compressed_dim', 'nrows')
            else:
                errmsg = errmsg.replace('compressed_indices', 'ccol_indices') \
                               .replace('plain_indices', 'row_indices') \
                               .replace('plain_dim', 'nrows') \
                               .replace('compressed_dim', 'ncols')

            if target == 'sparse_compressed_tensor_no_size' and label in {
                    'invalid size', 'invalid batchsize', 'invalid compressed_indices shape', 'invalid max(plain_indices)',
                    'invalid blocksize'}:
                # Skip invalid size input as a valid size is estimated for other inputs
                continue

            with self.assertRaisesRegex(RuntimeError, errmsg):
                if target == 'validate_sparse_compressed_tensor_args':
                    torch._validate_sparse_compressed_tensor_args(compressed_indices, plain_indices, values, size, layout)
                elif target == 'sparse_compressed_tensor':
                    torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, size, layout=layout)
                elif target == 'sparse_compressed_tensor_no_size':
                    torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, layout=layout)
                else:
                    raise NotImplementedError(target)

    @skipMeta
    @onlyCPU
    @largeTensorTest("30GB", "cpu")
    def test_invalid_input_csr_large(self):
        rows = 2 ** 31
        with self.assertRaisesRegex(RuntimeError, '32-bit integer overflow in row dimension'):
            torch.sparse_csr_tensor(torch.arange(rows + 1, dtype=torch.int32) // rows,
                                    torch.tensor([0], dtype=torch.int32),
                                    torch.tensor([1]), (rows, 1))
        torch.sparse_csr_tensor(torch.arange(rows + 1, dtype=torch.int64) // rows,
                                torch.tensor([0], dtype=torch.int64),
                                torch.tensor([1]), (rows, 1))

        cols = 2 ** 31
        with self.assertRaisesRegex(RuntimeError, '32-bit integer overflow in column dimension'):
            torch.sparse_csr_tensor(torch.arange(2, dtype=torch.int32),
                                    torch.tensor([0], dtype=torch.int32),
                                    torch.tensor([1]), (1, cols))
        torch.sparse_csr_tensor(torch.arange(2, dtype=torch.int64),
                                torch.tensor([0], dtype=torch.int64),
                                torch.tensor([1]), (1, cols))

        nnz = 2 ** 31
        with self.assertRaisesRegex(RuntimeError, '32-bit integer overflow in nnz'):
            # nnz cannot be stored in int32 crow_indices
            # but the `crow_indices[..., -1] == nnz`` check happens after the overflow validation
            # So we can use `nnz - 1` here to avoid `value cannot be converted to type int32 without overflow`
            # during construction of crow_indices
            torch.sparse_csr_tensor(torch.tensor([0, nnz // 2, nnz - 1], dtype=torch.int32),
                                    torch.arange(nnz // 2, dtype=torch.int32).repeat(2),
                                    torch.ones(nnz, dtype=torch.int8), (2, nnz // 2))
        torch.sparse_csr_tensor(torch.tensor([0, nnz // 2, nnz], dtype=torch.int64),
                                torch.arange(nnz // 2, dtype=torch.int64).repeat(2),
                                torch.ones(nnz, dtype=torch.int8), (2, nnz // 2))

    @skipMeta
    @onlyCPU
    @all_sparse_compressed_layouts()
    def test_dim(self, layout):
        for (compressed_indices, plain_indices, values), kwargs in self.generate_simple_inputs(layout, output_tensor=False):
            size = kwargs['size']
            batch_dim = compressed_indices.dim() - 1
            sparse_dim = 2
            block_dim = 2 if layout in {torch.sparse_bsr, torch.sparse_bsc} else 0
            dense_dim = values.dim() - batch_dim - block_dim - 1
            sparse = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, size, layout=layout)
            self.assertEqual(sparse.sparse_dim(), sparse_dim)
            self.assertEqual(sparse.dense_dim(), dense_dim)


    @skipMeta
    @all_sparse_compressed_layouts()
    @dtypes(*all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16))
    def test_to_dtype(self, layout, device, dtype):
        # to_dense does not support hybrid inputs
        for sparse in self.generate_simple_inputs(layout, dtype=dtype, device=device, enable_hybrid=False):
            for to_dtype in all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16):
                sparse_to_dtype = sparse.to(to_dtype)
                dense_to_dtype = sparse.to_dense().to(to_dtype)
                self.assertEqual(sparse_to_dtype.to_dense(), dense_to_dtype)

    @skipMeta
    @all_sparse_compressed_layouts()
    @dtypes(torch.double)
    def test_pickle(self, layout, dtype, device):
        import pickle

        for sparse in self.generate_simple_inputs(layout, device=device, dtype=dtype):
            serialized = pickle.dumps(sparse)
            sparse_loaded = pickle.loads(serialized)

            self.assertEqual(sparse, sparse_loaded)

    @all_sparse_compressed_layouts()
    @parametrize("index_dtype", [torch.int32, torch.int64])
    @dtypes(*all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool))
    def test_select_copy(self, device, dtype, index_dtype, layout):

        def is_view_of(base, other):
            # a shameless copy of TestViewOps.is_view_of
            if (
                not other._is_view() or
                other is base or
                other._base is not base or
                base.device != other.device
            ):
                return False
            if base.device.type in ('cpu', 'cuda'):
                if base.untyped_storage().data_ptr() != other.untyped_storage().data_ptr():
                    return False
            return True

        kwargs = dict(device=device, dtype=dtype, index_dtype=index_dtype)
        for sparse, dense in zip(self.generate_simple_inputs(layout, **kwargs),
                                 self.generate_simple_inputs(torch.strided, **kwargs)):
            if layout in {torch.sparse_csr, torch.sparse_bsr}:
                n_batchdim = sparse.crow_indices().ndim - 1
            elif layout in {torch.sparse_csc, torch.sparse_bsc}:
                n_batchdim = sparse.ccol_indices().ndim - 1
            else:
                assert 0  # unreachable
            self.assertEqual(sparse, dense)
            for dim in range(sparse.ndim):
                if sparse.shape[dim] == 0:
                    with self.assertRaisesRegex(IndexError, "index 0 out of range for tensor of size"):
                        torch.select_copy(sparse, dim, 0)
                    with self.assertRaisesRegex(IndexError, "index 0 out of range for tensor of size"):
                        torch.select_copy(dense, dim, 0)
                elif n_batchdim and dim >= n_batchdim and dim < n_batchdim + 2:
                    with self.assertRaisesRegex(
                            RuntimeError,
                            "selecting sparse dimensions is not supported for batched sparse compressed tensors"):
                        torch.select_copy(sparse, dim, 0)
                else:
                    for index in {0, sparse.shape[dim] // 2, sparse.shape[dim] - 1}:
                        dense_select = torch.select_copy(dense, dim, index)
                        sparse_select = torch.select_copy(sparse, dim, index)
                        self.assertEqual(sparse_select, dense_select)
                        self.assertFalse(is_view_of(sparse_select.values(), sparse.values()))


def _npref_block_addmm_addmv(c, a, b, alpha, beta):
    return alpha * (a @ b) + beta * c


class TestSparseCSR(TestCase):

    def test_csr_stride(self):
        a = self.genSparseCSRTensor((3, 3), 3, dtype=torch.float, device=self.device_type, index_dtype=torch.int64)

        with self.assertRaisesRegex(RuntimeError, "Sparse CSR tensors do not have strides"):
            a.stride()

        with self.assertRaisesRegex(RuntimeError, "Sparse CSR tensors do not have strides"):
            a.stride(-1)

    def test_csr_storage(self):
        a = self.genSparseCSRTensor((3, 3), 3, dtype=torch.float, device=self.device_type, index_dtype=torch.int64)

        with self.assertRaisesRegex(RuntimeError, "Cannot access storage of SparseCsrTensorImpl"):
            a.storage()

    def test_csr_is_contiguous(self):
        a = self.genSparseCSRTensor((3, 3), 3, dtype=torch.float, device=self.device_type, index_dtype=torch.int64)

        with self.assertRaisesRegex(RuntimeError, "Sparse CSR tensors do not have is_contiguous"):
            a.is_contiguous()

    @onlyCPU
    @largeTensorTest("20GB", "cpu")
    def test_csr_nnz(self):
        # Tests the limits of the number of specified elements in CSR tensors, see gh-102520.
        for nnz in [0, 2**31]:
            rows, cols = 1, max(nnz, 1)
            crow_indices = torch.tensor([0, nnz], dtype=torch.int64)
            col_indices = torch.arange(nnz, dtype=torch.int64)
            values = torch.ones(nnz, dtype=torch.int8)
            a = torch.sparse_csr_tensor(crow_indices, col_indices, values, (rows, cols))
            self.assertEqual(a._nnz(), nnz)

    def test_csr_double_to_sparse_csr(self):
        a = self.genSparseCSRTensor((3, 3), 3, dtype=torch.float, device=self.device_type, index_dtype=torch.int64)
        a.to_sparse_csr().to_sparse_csr()

    @all_sparse_compressed_layouts()
    @parametrize("index_dtype", [torch.int32, torch.int64])
    @dtypes(*all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool))
    def test_select(self, device, dtype, index_dtype, layout):
        compressed_indices_mth = {
            torch.sparse_csr: torch.Tensor.crow_indices,
            torch.sparse_bsr: torch.Tensor.crow_indices,
            torch.sparse_csc: torch.Tensor.ccol_indices,
            torch.sparse_bsc: torch.Tensor.ccol_indices,
        }[layout]

        plain_indices_mth = {
            torch.sparse_csr: torch.Tensor.col_indices,
            torch.sparse_bsr: torch.Tensor.col_indices,
            torch.sparse_csc: torch.Tensor.row_indices,
            torch.sparse_bsc: torch.Tensor.row_indices,
        }[layout]
        create_tensor_mth = {
            torch.sparse_csr: torch.sparse_csr_tensor,
            torch.sparse_bsr: torch.sparse_bsr_tensor,
            torch.sparse_csc: torch.sparse_csc_tensor,
            torch.sparse_bsc: torch.sparse_bsc_tensor,
        }[layout]

        shape = (2, 3, 6, 10)
        nnz = 6
        blocksize = (2, 2) if layout in {torch.sparse_bsr, torch.sparse_bsc} else ()
        sparse = self.genSparseCompressedTensor(
            shape, nnz, device=device, layout=layout, dtype=dtype, index_dtype=index_dtype, blocksize=blocksize)
        comp_indices = compressed_indices_mth(sparse)
        plain_indices = plain_indices_mth(sparse)
        values = sparse.values()

        # select from batch dimensions
        sparse_selected12 = sparse.select(1, 2)
        expected_sparse_selected12 = create_tensor_mth(comp_indices.select(1, 2).contiguous(),
                                                       plain_indices.select(1, 2).contiguous(),
                                                       values.select(1, 2).contiguous(),
                                                       size=(2, 6, 10),
                                                       dtype=dtype,
                                                       device=device)
        self.assertEqual(expected_sparse_selected12, sparse_selected12)

        # selecting rows/col with batch dims not allowed
        sparse_non_batched = sparse[0, 0]
        # select from sparse dimensions
        for select_args in [(0, 0), (1, 1)]:
            sparse_selected = sparse_non_batched.select(*select_args)
            dense_selected = sparse_non_batched.to_dense().select(*select_args)
            self.assertEqual(dense_selected, sparse_selected)

        self.assertEqual(sparse[0, 0, 0, 0], sparse.to_dense()[0, 0, 0, 0])
        # assigning to sparse through indexing is disabled
        with self.assertRaisesRegex(TypeError, "Cannot assign to a sparse tensor"):
            sparse[0, 0, 0, 0] = 99.0

        # select from sparse dimensions without removing batch dims
        msg = "selecting sparse dimensions is not supported for batched sparse compressed tensors."
        with self.assertRaisesRegex(RuntimeError, msg):
            sparse.select(-2, 0)

        with self.assertRaisesRegex(RuntimeError, msg):
            sparse.select(-1, 0)

    @skipMeta
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_resize(self, device, dtype):

        def numel(tensor):
            r = 1
            for s in tensor.shape:
                r *= s
            return r

        batch_shapes = [(), (2,), (2, 3)]
        for index_dtype, b in zip([torch.int32, torch.int64], batch_shapes):
            shape = (*b, 2, 3)
            nnz = 6
            a = self.genSparseCSRTensor(shape, nnz, dtype=dtype, device=device, index_dtype=index_dtype)
            self.assertEqual(a.numel(), numel(a))

            new_shape = (*b, 4, 5)
            a.resize_(new_shape)

            self.assertEqual(a.shape, new_shape)
            # resize to larger shape doesn't add specified elements
            self.assertEqual(a._nnz(), nnz)
            self.assertEqual(a.numel(), numel(a))

            new_shape = (*b, 1, 5)
            a.resize_(new_shape)

            self.assertEqual(a.shape, new_shape)
            # resize to smaller shape trims specified elements
            self.assertEqual(a._nnz(), 5)
            self.assertEqual(a.numel(), numel(a))

            # trim batched dimensions
            a.resize_(new_shape[-2], new_shape[-1])
            self.assertEqual(a.shape, (new_shape[-2], new_shape[-1]))
            self.assertEqual(a._nnz(), 5)
            self.assertEqual(a.numel(), numel(a))

    @skipMeta
    @dtypes(torch.float, torch.bool)
    @all_sparse_compressed_layouts()
    def test_resize_as_sparse_compressed(self, device, dtype, layout):

        def _check_resize_b_as_a(b, a):
            br = b.clone()
            br.resize_as_sparse_(a)

            # shape is inherited from a
            self.assertEqual(a.shape, br.shape)
            # other metadata is not affected
            self.assertEqual(b.layout, br.layout)
            self.assertEqual(b.device, br.device)
            self.assertEqual(b.dtype, br.dtype)

            def _get_compressed_plain_inds(t):
                compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[t.layout]
                return compressed_indices_mth(t), plain_indices_mth(t)

            br_compressed_indices, br_plain_indices = _get_compressed_plain_inds(br)
            br_values = br.values()

            b_compressed_indices, b_plain_indices = _get_compressed_plain_inds(b)
            a_compressed_indices, a_plain_indices = _get_compressed_plain_inds(a)
            self.assertEqual(a_plain_indices.shape, br_plain_indices.shape)
            self.assertEqual(a_compressed_indices.shape, br_compressed_indices.shape)
            # We don't check the content of br_plain_indices and br_compressed_indices
            # because it is not well-defined (the content depends on the original
            # shape of `b` that `resize_as` ought to discard) nor needed (the
            # subsequent operation likely updates the indices and values of `b` anyway).
            # the device/dtype of indices should always be unaffected
            self.assertEqual(b_plain_indices.dtype, br_plain_indices.dtype)
            self.assertEqual(b_plain_indices.device, br_plain_indices.device)
            self.assertEqual(b_compressed_indices.dtype, br_compressed_indices.dtype)
            self.assertEqual(b_compressed_indices.device, br_compressed_indices.device)
            # values are generated empty, shape is updated
            self.assertEqual(a.values().shape, br_values.shape)
            # the device/dtype of indices should always be unaffected
            b_values = b.values()
            self.assertEqual(b_values.dtype, br_values.dtype)
            self.assertEqual(b_values.device, br_values.device)
            # nnz will be picked up from a via new shape of values
            self.assertEqual(a._nnz(), br._nnz())

            # post resize the invariants of the layout are respected
            torch._validate_sparse_compressed_tensor_args(br_compressed_indices, br_plain_indices, br_values, br.shape,
                                                          br.layout)

        block_sparse = layout in (torch.sparse_bsr, torch.sparse_bsc)
        shape = (2, 1, 6, 4)
        nnz = 4
        blocksize = (2, 1) if block_sparse else ()
        for index_dtype in [torch.int32, torch.int64]:
            a = self.genSparseCompressedTensor(shape,
                                               layout=layout,
                                               device=device,
                                               index_dtype=index_dtype,
                                               dtype=dtype,
                                               nnz=nnz,
                                               blocksize=blocksize)

            # same size, resize should not trigger
            b = self.genSparseCompressedTensor(shape,
                                               layout=layout,
                                               device=device,
                                               index_dtype=index_dtype,
                                               dtype=dtype,
                                               nnz=nnz,
                                               blocksize=blocksize)

            # This test will not always trigger a resize, if the layouts are the same nothing should happen to b.
            # The invariants of the function as checked should still hold
            _check_resize_b_as_a(b, a)

            # same ndim, but bigger, more nnz, different dtype, different blocksize if blocked
            b = self.genSparseCompressedTensor(tuple(s * 2 for s in shape),
                                               layout=layout,
                                               device=device,
                                               dtype=torch.chalf,
                                               index_dtype=torch.int64 if index_dtype == torch.int32 else torch.int32,
                                               nnz=nnz * 2,
                                               blocksize=tuple(2 * bi for bi in blocksize))
            _check_resize_b_as_a(b, a)

            # different device, only check on cuda pass as we know we are testing in an environment
            # that has multiple devices

            # TODO: .cpu() does not seem to work correctly for sparse. Causes a call to `copy_` which
            # complains about incompatible nnz between src and self?
            if torch.device(device).type == 'cuda' and (layout not in (torch.sparse_bsc, torch.sparse_bsr)):
                a_cpu = self.genSparseCompressedTensor(shape,
                                                       layout=layout,
                                                       device='cpu',
                                                       index_dtype=index_dtype,
                                                       dtype=dtype,
                                                       nnz=nnz,
                                                       blocksize=blocksize)
                _check_resize_b_as_a(b, a)

            # error on a strided
            a_strided = a.to_dense()
            with self.assertRaisesRegex(
                    RuntimeError, r'resize_as_sparse_compressed_: src  expected sparse compressed tensor layout'):
                b.resize_as_sparse_(a_strided)

            # error on b strided
            b_strided = b.to_dense()
            with self.assertRaisesRegex(
                    RuntimeError, r'resize_as_sparse_compressed_: self  expected sparse compressed tensor layout'):
                b_strided.resize_as_sparse_(a)

            # error if layout does not match, transpose induces layout flip
            with self.assertRaisesRegex(RuntimeError,
                                        r"resize_as_sparse_compressed_tensor_: self and src must have the same layout"):
                b.transpose(-2, -1).resize_as_sparse_(a)
            with self.assertRaisesRegex(RuntimeError,
                                        r"resize_as_sparse_compressed_tensor_: self and src must have the same layout"):
                b.resize_as_sparse_(a.transpose(-2, -1))

    @skipMeta
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_resize_errors(self, device, dtype):
        for index_dtype in [torch.int32, torch.int64]:
            shape = (2, 3)
            nnz = 6
            a = self.genSparseCSRTensor(shape, nnz, dtype=dtype, device=device, index_dtype=index_dtype)

            with self.assertRaisesRegex(RuntimeError, "torch.resize_: Only batched sparse CSR matrices are supported"):
                new_shape = (4,)
                a.resize_(new_shape)

            # resizing of columns to smaller size is not implemented
            with self.assertRaisesRegex(
                RuntimeError,
                "torch.resize_: Resizing columns of sparse CSR tensors to a smaller value is not supported.",
            ):
                new_shape = (2, 2)
                a.resize_(new_shape)

    @skipIfTorchDynamo()
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_sparse_csr_from_dense(self, device, dtype):
        dense = torch.tensor([[4, 5, 0], [0, 0, 0], [1, 0, 0]], dtype=dtype, device=device)
        sparse = dense.to_sparse_csr()
        self.assertEqual(torch.tensor([0, 2, 2, 3], dtype=torch.int64), sparse.crow_indices())
        self.assertEqual(torch.tensor([0, 1, 0], dtype=torch.int64), sparse.col_indices())
        self.assertEqual(torch.tensor([4, 5, 1], dtype=dtype), sparse.values())

        dense = torch.tensor([[0, 0, 0], [0, 0, 1], [1, 0, 0]], dtype=dtype, device=device)
        sparse = dense.to_sparse_csr()
        self.assertEqual(torch.tensor([0, 0, 1, 2], dtype=torch.int64), sparse.crow_indices())
        self.assertEqual(torch.tensor([2, 0], dtype=torch.int64), sparse.col_indices())
        self.assertEqual(torch.tensor([1, 1], dtype=dtype), sparse.values())

        dense = torch.tensor([[2, 2, 2], [2, 2, 2], [2, 2, 2]], dtype=dtype, device=device)
        sparse = dense.to_sparse_csr()
        self.assertEqual(torch.tensor([0, 3, 6, 9], dtype=torch.int64), sparse.crow_indices())
        self.assertEqual(torch.tensor([0, 1, 2] * 3, dtype=torch.int64), sparse.col_indices())
        self.assertEqual(torch.tensor([2] * 9, dtype=dtype), sparse.values())

    def _test_sparse_compressed_to_dense(self, device, dtype, layout):
        compressed_format_str = str(layout)[-3:]

        def to_compressed(t):
            return getattr(t, f"to_sparse_{compressed_format_str}")()

        def compressed_constructor(*input, **kwargs):
            constructor = getattr(torch, f"sparse_{compressed_format_str}_tensor")
            return constructor(*input, **kwargs)

        def get_dense_shape(shape, batch_ndim):
            if layout is torch.sparse_csc:
                compressed_dims_slice = slice(batch_ndim + 1, batch_ndim - 1, -1)
            else:
                compressed_dims_slice = slice(batch_ndim, batch_ndim + 2)
            return shape[:batch_ndim] + shape[compressed_dims_slice] + shape[batch_ndim + 2:]

        def transpose(t, batch_ndim):
            if layout is torch.sparse_csc:
                return t.transpose(batch_ndim, batch_ndim + 1)
            return t

        mn = [5, 2, 0]
        for (m, n) in itertools.product(mn, mn):
            size = (m, n)
            dense = make_tensor(size, dtype=dtype, device=device)
            sparse = to_compressed(dense)
            self.assertEqual(sparse.to_dense(), dense)

        batch_shape = (2, 3)
        compressed_indices = torch.tensor([0, 3, 5], device=device).repeat(6, 1).reshape(*batch_shape, -1)
        plain_indices = torch.tensor([0, 1, 2, 0, 1], device=device).repeat(6, 1).reshape(*batch_shape, -1)
        values = torch.tensor([1, 2, 1, 3, 4], device=device, dtype=dtype).repeat(6, 1).reshape(*batch_shape, -1)
        sparse = compressed_constructor(compressed_indices, plain_indices, values, dtype=dtype, device=device)
        dense_shape = get_dense_shape(sparse.shape, len(batch_shape))
        dense = torch.tensor([[1, 2, 1], [3, 4, 0]], dtype=dtype, device=device).repeat(6, 1).reshape(dense_shape)
        self.assertEqual(sparse.to_dense(), transpose(dense, len(batch_shape)))

    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_sparse_csr_to_dense(self, device, dtype):
        self._test_sparse_compressed_to_dense(device, dtype, torch.sparse_csr)

    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_sparse_csc_to_dense(self, device, dtype):
        self._test_sparse_compressed_to_dense(device, dtype, torch.sparse_csc)

    @skipMeta
    @skipCPUIfNoMklSparse
    @coalescedonoff
    @dtypes(torch.double)
    def test_coo_to_csr_convert(self, device, dtype, coalesced):
        with self.assertRaisesRegex(RuntimeError, "Input is supposed to be a vector"):
            torch._convert_indices_from_coo_to_csr(
                torch.randint(100, (5, 5), device=device),
                size=100)

        size = (5, 5)
        sparse_dim = 2
        nnz = 10
        sparse_coo, _, _ = self.genSparseTensor(size, sparse_dim, nnz, coalesced, device, dtype)
        sparse_csr = sparse_coo.to_sparse_csr()

        self.assertTrue(sparse_csr.is_sparse_csr)
        self.assertEqual(sparse_csr.to_dense(), sparse_coo.to_dense())

        vec = torch.randn((5, 1), dtype=dtype, device=device)
        coo_product = sparse_coo.matmul(vec)
        csr_product = sparse_csr.matmul(vec)

        self.assertEqual(coo_product, csr_product)

        vec = torch.randn((100, 1), dtype=dtype, device=device)
        index = torch.tensor([
            [1, 0, 35, 14, 39, 6, 71, 66, 40, 27],
            [92, 31, 62, 50, 22, 65, 89, 74, 56, 34],
        ], dtype=torch.int32)
        values = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=dtype, device=device)
        coo = torch.sparse_coo_tensor(index, values, torch.Size([100, 100]), dtype=dtype, device=device)
        csr = coo.to_sparse_csr()

        self.assertEqual(coo.matmul(vec), csr.matmul(vec))

        col_indices = torch.tensor([
            31, 92, 65, 50, 34, 62, 22, 56, 74, 89
        ], dtype=torch.int64, device=device)
        self.assertEqual(csr.col_indices(), col_indices)

        values = torch.tensor([2, 1, 6, 4, 10, 3, 5, 9, 8, 7], dtype=dtype, device=device)
        self.assertEqual(csr.values(), values)

    @parametrize("blocksize", [2, 4])
    @dtypes((torch.double, torch.int32), (torch.double, torch.int64))
    @unittest.skipIf(not TEST_SCIPY, "SciPy not found")
    @skipMeta
    def test_csr_to_block_csr(self, device, dtypes, blocksize):
        for shape in [(24, 24), (12, 24)]:
            dtype, index_dtype = dtypes
            m, k = shape
            nnz = random.randint(0, m * k)
            t = self.genSparseCSRTensor((m * blocksize, k * blocksize), nnz, dtype=dtype,
                                        device=device, index_dtype=index_dtype)
            st = sp.csr_matrix((t.values().cpu(), t.col_indices().cpu(), t.crow_indices().cpu()), shape=tuple(t.size()))
            block_t = t.to_sparse_bsr((blocksize, blocksize))
            self.assertEqual(block_t.values().dim(), 3)
            self.assertTrue(block_t.layout == torch.sparse_bsr)
            block_st = st.tobsr(blocksize=(blocksize, blocksize))
            block_st.sort_indices()
            self.assertEqual(block_t.values().cpu(), block_st.data)
            self.assertEqual(block_t.col_indices().cpu(), torch.tensor(block_st.indices).to(index_dtype))
            self.assertEqual(block_t.crow_indices().cpu(), torch.tensor(block_st.indptr).to(index_dtype))

    @dtypes(torch.double)
    @unittest.skipIf(not TEST_SCIPY, "SciPy not found")
    def test_csr_to_block_csr_errors(self, device, dtype):
        for index_dtype in [torch.int32, torch.int64]:
            nnz = 15
            t = self.genSparseCSRTensor((16, 16), nnz, dtype=dtype,
                                        device=device, index_dtype=index_dtype)

            with self.assertRaisesRegex(RuntimeError,
                                        r"tensor sparse size \(.*,.*\) must be divisible by given blocksize \(.*,.*\)"):
                block_t = t.to_sparse_bsr((5, 5))

    # TODO: Support auto generation of device check for sparse tensors
    # See: https://github.com/pytorch/pytorch/issues/59058
    @onlyCUDA
    @dtypes(torch.double)
    def test_matmul_device_mismatch(self, device, dtype):
        cpu = torch.rand((10, 10))
        cuda = cpu.cuda()
        for s, m1, m2 in itertools.product((cpu, cuda), repeat=3):
            csr = m1.to_sparse()
            if s.device == csr.device == m2.device:
                torch.addmm(s, csr, m2)
            else:
                with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"):
                    torch.addmm(s, csr, m2)

    @skipCPUIfNoMklSparse
    @skipCUDAIfNoSparseGeneric
    @dtypes(*floating_and_complex_types())
    @dtypesIfCUDA(*floating_and_complex_types_and(
                  *[torch.half] if SM53OrLater else [],
                  *[torch.bfloat16] if SM80OrLater else []))
    def test_csr_matvec(self, device, dtype):

        if TEST_WITH_ROCM and (dtype == torch.half or dtype == torch.bfloat16):
            self.skipTest("ROCm doesn't work with half dtypes correctly.")

        side = 100
        for index_dtype in [torch.int32, torch.int64]:
            csr = self.genSparseCSRTensor((side, side), 1000, device=device, dtype=dtype, index_dtype=index_dtype)
            vec = torch.randn(side, dtype=dtype, device=device)

            res = csr.matmul(vec)
            expected = csr.to_dense().matmul(vec)

            self.assertEqual(res, expected)

            bad_vec = torch.randn(side + 10, dtype=dtype, device=device)
            err_msg = "size mismatch, got"
            with self.assertRaisesRegex(RuntimeError, err_msg):
                csr.matmul(bad_vec)

    @onlyCUDA
    # hmm, the test passes ok on CUDA when Rocm is not available:
    @skipCUDAIfRocmVersionLessThan((5, 2))
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    def test_baddbmm(self, device, dtype):

        # TODO: disable the invariant checks within torch.baddbmm that
        # constructs unconventional csr tensors leading to
        # RuntimeError: tensor dimensionality must be sum of batch,
        #     base, and dense dimensionalities (=0 + 2 + 0) but got 3
        # when invariant checking is enabled. When done, undecorate run_test.
        @torch.sparse.check_sparse_tensor_invariants(enable=False)
        def run_test(c, a, a_batched, b, op_b=False, op_out=False, *, dtype=None, device=None):
            alpha = complex(random.random(), random.random()) if dtype.is_complex else random.random()
            beta = complex(random.random(), random.random()) if dtype.is_complex else random.random()
            b = b.mH if (op_b and a.shape == b.shape) else b

            actual = torch.baddbmm(c, a_batched, b, alpha=alpha, beta=beta)

            out = torch.empty_like(c.mH if op_out and a.shape == b.shape else c)
            torch.baddbmm(c, a_batched, b, alpha=alpha, beta=beta, out=out)

            expected = [torch.addmm(c[i], a, b[i], alpha=alpha, beta=beta) for i in range(c.shape[0])]
            expected = torch.stack(expected, 0)

            self.assertEqual(actual, out)
            self.assertEqual(actual, expected)

        for index_dtype in [torch.int32, torch.int64]:
            for (m, n, k), batch_size, noncontiguous in zip(itertools.product([2, 5], repeat=3), [1, 3], [True, False]):
                nnz = random.randint(0, m * k)
                a = self.genSparseCSRTensor((m, k), nnz, dtype=dtype, device=device, index_dtype=index_dtype)

                # a_batched is a regular CSR tensor but with a batch dimension in the shape
                a_batched = torch.sparse_csr_tensor(
                    a.crow_indices(), a.col_indices(), a.values(), (batch_size, m, k), check_invariants=False)

                b = make_tensor((batch_size, k, n), dtype=dtype, device=device, noncontiguous=noncontiguous)
                c = make_tensor((batch_size, m, n), dtype=dtype, device=device, noncontiguous=noncontiguous)
                for op_b, op_out in itertools.product([True, False], repeat=2):
                    run_test(c, a, a_batched, b, op_b, op_out, dtype=dtype, device=device)

    @onlyCUDA
    @unittest.skipIf(TEST_WITH_ROCM, "Only CUDA 11+ is supported")
    @skipCUDAIfNoSparseGeneric
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    def test_bmm(self, device, dtype):
        def run_test(a, a_batched, b, op_b=False, op_out=False, *, dtype=None, device=None):
            b = b.mH if (op_b and a.shape == b.shape) else b

            actual = torch.bmm(a_batched, b)

            out = torch.empty_like(actual.mH if op_out and a.shape == b.shape else actual)
            torch.bmm(a_batched, b, out=out)

            expected = [torch.mm(a, b[i]) for i in range(b.shape[0])]
            expected = torch.stack(expected, 0)

            self.assertEqual(actual, out)
            self.assertEqual(actual, expected)

        for index_dtype in [torch.int32, torch.int64]:
            for (m, n, k), batch_size, noncontiguous in zip(itertools.product([2, 5], repeat=3), [1, 3], [True, False]):
                nnz = random.randint(0, m * k)
                a = self.genSparseCSRTensor((m, k), nnz, dtype=dtype, device=device, index_dtype=index_dtype)

                # a_batched is a regular CSR tensor but with a batch
                # dimension in the shape. It is unorthodox in PyTorch
                # to represent a batch sparse tensor in this way,
                # hence checking the tensor invariants is locally
                # turned off.
                a_batched = torch.sparse_csr_tensor(
                    a.crow_indices(), a.col_indices(), a.values(), (batch_size, m, k), check_invariants=False)

                b = make_tensor((batch_size, k, n), dtype=dtype, device=device, noncontiguous=noncontiguous)
                for op_b, op_out in itertools.product([True, False], repeat=2):
                    run_test(a, a_batched, b, op_b, op_out, dtype=dtype, device=device)

    def run_test_block_addmm_addmv(self,
                                   addmv_addmm,
                                   c,
                                   a,
                                   b,
                                   op_b=False,
                                   op_out=False,
                                   *,
                                   dtype=None,
                                   device=None,
                                   ref=_npref_block_addmm_addmv):
        alpha = complex(random.random(), random.random()) if dtype.is_complex else random.random()
        beta = complex(random.random(), random.random()) if dtype.is_complex else random.random()
        b = b.mH if (op_b and a.shape == b.shape) else b

        actual = addmv_addmm(c, a, b, alpha=alpha, beta=beta)

        out = torch.empty_like(c.mH if op_out and a.shape == b.shape else c)
        addmv_addmm(c, a, b, alpha=alpha, beta=beta, out=out)
        expected = ref(c, a, b, alpha, beta)

        self.assertEqual(actual, out)
        self.assertEqual(actual, expected, lambda msg: f"{msg}\na={a}\nc={c}\nb={b}\nalpha={alpha} beta={beta}")

    # TODO: block_size 1 is broken
    @parametrize("block_size", [2, 3])
    @parametrize("index_dtype", [torch.int32, torch.int64])
    @parametrize("noncontiguous", [True, False])
    @skipCPUIfNoMklSparse
    @unittest.skipIf(not TEST_SCIPY, "SciPy not found")
    @skipIfTorchDynamo("raises 'sparse matrix length is ambiguous; use getnnz()'")
    @dtypes(*floating_and_complex_types())
    @dtypesIfCUDA(*floating_and_complex_types_and(
                  *[torch.half] if SM53OrLater else [],
                  *[torch.bfloat16] if SM80OrLater else []))
    @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
                        torch.float64: 1e-5, torch.complex128: 1e-5,
                        torch.float16: 1e-3, torch.bfloat16: 1e-3})
    def test_block_addmm(self, device, dtype, index_dtype, block_size, noncontiguous):

        def make_transposed_addmm_op(f):

            def tt(t):
                if isinstance(t, torch.Tensor):
                    return t.transpose(-2, -1)
                else:
                    # assume numpy/scipy spmatrix
                    return t.transpose()

            @functools.wraps(f)
            def wrapper(c, a, b, alpha=None, beta=None, out=None):
                if out is not None:
                    # the ref takes no out kwarg
                    assert isinstance(out, torch.Tensor)
                    # transpose inplace to propagate out to checking context
                    out.transpose_(-2, -1)
                    return f(tt(c), tt(b), tt(a), alpha=alpha, beta=beta, out=out)
                else:
                    return f(tt(c), tt(b), tt(a), alpha=alpha, beta=beta)

            return wrapper

        def ref_sp_numpy(c, a, b, alpha=None, beta=None, out=None):

            def prep_input(t):

                def to_sp_block_compressed(t):

                    if t.layout is torch.sparse_bsc:
                        tt = t.transpose(-1, -2)
                    else:
                        tt = t

                    t_sp_bsr = sp.bsr_matrix(
                        (
                            tt.values().cpu().numpy(),
                            tt.col_indices().cpu().numpy(),
                            tt.crow_indices().cpu().numpy(),
                        ),
                        shape=tt.shape,
                    )

                    if t.layout is torch.sparse_bsc:
                        return t_sp_bsr.transpose()
                    else:
                        return t_sp_bsr

                if t.layout is not torch.strided:
                    return to_sp_block_compressed(t)
                else:
                    return t.cpu().resolve_conj().numpy()

            res = _npref_block_addmm_addmv(
                *(prep_input(t) for t in (c, a, b)),
                alpha,
                beta
            )

            if out is not None:
                out.copy_(res)
                return out
            else:
                return res

        def ref_half_bfloat16(c, a, b, alpha=None, beta=None, out=None):
            res = alpha * (a.to_dense().to(torch.float32) @ b.to_dense().to(torch.float32)).to(a.dtype) + beta * c
            if out is not None:
                out.copy_(res)
                return out
            else:
                return res

        if dtype in (torch.half, torch.bfloat16):
            ref = ref_half_bfloat16
        else:
            ref = ref_sp_numpy

        for (m, n, k) in itertools.product([2, 5], repeat=3):
            nnz = random.randint(0, m * k)
            a = self.genSparseCSRTensor((m, k), nnz, dtype=dtype, device=device, index_dtype=index_dtype)
            a_data = make_tensor((nnz, block_size, block_size), dtype=dtype, device=device)
            a_data = a_data.mT if noncontiguous else a_data
            a = torch.sparse_bsr_tensor(a.crow_indices(), a.col_indices(),
                                        a_data, (m * block_size, k * block_size), check_invariants=False)
            b = make_tensor((k * block_size, n * block_size), dtype=dtype, device=device, noncontiguous=noncontiguous)
            c = make_tensor((m * block_size, n * block_size), dtype=dtype, device=device, noncontiguous=noncontiguous)
            for op_b, op_out in itertools.product([True, False], repeat=2):
                self.run_test_block_addmm_addmv(torch.addmm, c, a, b, op_b, op_out, dtype=dtype, device=device, ref=ref)
                self.run_test_block_addmm_addmv(make_transposed_addmm_op(torch.addmm),
                                                c,
                                                a,
                                                b,
                                                op_b,
                                                op_out,
                                                dtype=dtype,
                                                device=device,
                                                ref=make_transposed_addmm_op(ref))

    @parametrize("block_size", [2, 3])
    @parametrize("index_dtype", [torch.int32, torch.int64])
    @parametrize("noncontiguous", [True, False])
    @unittest.skipIf(not TEST_SCIPY, "SciPy not found")
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    def test_block_addmv(self, device, dtype, index_dtype, block_size, noncontiguous):
        # TODO: Explicitly disable block size 1 support
        # if (TEST_WITH_ROCM or not TEST_CUSPARSE_GENERIC) and block_size == 1:
        #     return
        def ref_block_addmv(c, a, b, alpha, beta):
            return _npref_block_addmm_addmv(c, a.to_dense(), b, alpha, beta)

        for (m, k) in itertools.product([2, 5], repeat=2):
            nnz = random.randint(0, m * k)
            if not noncontiguous:
                a = self.genSparseCSRTensor((m * block_size, k * block_size), nnz,
                                            dtype=dtype, device=device, index_dtype=index_dtype)
                a = a.to_sparse_bsr((block_size, block_size))
            else:
                a = self.genSparseCSRTensor((m, k), nnz, dtype=dtype, device=device, index_dtype=index_dtype)
                a_data = make_tensor((nnz, block_size, block_size), dtype=dtype, device=device)
                a_data = a_data.mT if noncontiguous else a_data   # Test column-major blocks
                a = torch.sparse_bsr_tensor(a.crow_indices(), a.col_indices(),
                                            a_data, (m * block_size, k * block_size), check_invariants=False)
            b = make_tensor((k * block_size,), dtype=dtype, device=device, noncontiguous=noncontiguous)
            c = make_tensor((m * block_size,), dtype=dtype, device=device, noncontiguous=noncontiguous)
            self.run_test_block_addmm_addmv(torch.addmv, c, a, b, dtype=dtype, device=device, ref=ref_block_addmv)

    @parametrize("matrix_shape", [(3, 3), (5, 7), (11, 9)], name_fn=lambda x: "shape_{}x{}".format(*x))
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    @onlyCPU
    def test_addmv(self, device, dtype, matrix_shape):
        mat = torch.randn(matrix_shape, dtype=dtype, device=device)
        mat[mat.real < 0] = 0
        sparse_mat = mat.to_sparse_csr()
        mvec = torch.randn((mat.size(1),), dtype=dtype, device=device)
        avec = torch.randn((mat.size(0),), dtype=torch.float64, device=device)
        ref_output = torch.addmv(avec, mat, mvec)
        output = torch.addmv(avec, sparse_mat, mvec)
        self.assertEqual(ref_output, output)

    @parametrize("block_size", [2, 3])
    @parametrize("index_dtype", [torch.int32, torch.int64])
    @parametrize("noncontiguous", [True, False])
    @skipCPUIfNoMklSparse
    @unittest.skipIf(not TEST_SCIPY, "SciPy not found")
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    def test_block_triangular_solve(self, device, dtype, index_dtype, block_size, noncontiguous):
        def run_test(a, b, upper, transpose, unitriangular, op_out):
            if unitriangular and self.device_type == 'cpu':
                # TODO: When unitriangular=True results are not correct on CPU
                return

            if not upper and self.device_type == 'cpu':
                # TODO: When upper=False some generated inputs might crash on CPU
                return

            actual = torch.triangular_solve(b, a, upper=upper, unitriangular=unitriangular, transpose=transpose)
            actual_X = actual.solution
            actual_A_clone = actual.cloned_coefficient
            self.assertTrue(actual_A_clone.numel() == 0)
            if a._nnz() == 0:
                self.assertTrue(actual_X.isnan().all())
                return

            # TODO: replace with torch method when implemented to_dense() on block sparse tensor
            a_bsr = sp.bsr_matrix(
                (
                    a.values().cpu().numpy(),
                    a.col_indices().cpu().numpy(),
                    a.crow_indices().cpu().numpy(),
                ),
                shape=a.shape,
            )
            expected_X, _ = torch.triangular_solve(
                b,
                torch.tensor(a_bsr.todense(), device=device),
                transpose=transpose,
                upper=upper,
                unitriangular=unitriangular)

            if expected_X.isnan().any():
                # TODO: zeros on the diagonal are not handled for CPU path
                # there's no way to query this info from MKL
                if self.device_type == 'cuda' and not TEST_WITH_ROCM:
                    self.assertTrue(actual_X.isnan().any() or actual_X.isinf().any())
                return

            self.assertEqual(actual_X, expected_X)

            out = torch.empty_like(b.mH if op_out and a.shape == b.shape else b)
            torch.triangular_solve(
                b, a,
                upper=upper, unitriangular=unitriangular, transpose=transpose, out=(out, actual_A_clone)
            )
            self.assertEqual(out, actual_X)
            self.assertEqual(out, expected_X)

        for (m, k) in itertools.product([2, 3], [1, 3]):
            nnz = random.randint(0, m * m)
            if not noncontiguous:
                a = self.genSparseCSRTensor((m * block_size, m * block_size), nnz,
                                            dtype=dtype, device=device, index_dtype=index_dtype)
                a = a.to_sparse_bsr((block_size, block_size))
            else:
                a = self.genSparseCSRTensor((m, m), nnz, dtype=dtype, device=device, index_dtype=index_dtype)
                a_data = make_tensor((nnz, block_size, block_size), dtype=dtype, device=device)
                a_data = a_data.mT if noncontiguous else a_data  # Test column-major blocks
                a = torch.sparse_bsr_tensor(a.crow_indices(), a.col_indices(),
                                            a_data, (m * block_size, m * block_size), check_invariants=False)
            b = make_tensor((m * block_size, k), dtype=dtype, device=device, noncontiguous=noncontiguous)

            for (upper, unitriangular, transpose, op_out) in itertools.product([True, False], repeat=4):
                run_test(a, b, upper, unitriangular, transpose, op_out)

    @skipCPUIfNoMklSparse
    @unittest.skipIf(TEST_WITH_ROCM, "Only CUDA 11+ is supported")
    @dtypes(torch.double)
    def test_mm(self, device, dtype):
        def test_shape(di, dj, dk, nnz0=None, nnz1=None):
            for index_dtype in [torch.int32, torch.int64]:
                alpha = random.random()
                beta = random.random()

                def _test_addmm(t, x, y):
                    # TODO: addmm doesn't support strided result for sparse inputs.
                    # res = beta * t  + alpha * (x @ y)
                    res = torch.addmm(t, x, y, beta=beta, alpha=alpha)
                    expected = torch.addmm(t, x.to_dense(), y.to_dense(), beta=beta, alpha=alpha)
                    self.assertEqual(res, expected)

                    res = torch.addmm(t, x, y)
                    expected = torch.addmm(t, x.to_dense(), y.to_dense())
                    self.assertEqual(res, expected)

                def _test_mm(x, y):
                    res = torch.mm(x, y)
                    expected = torch.mm(x.to_dense(), y.to_dense())
                    if x.layout is torch.strided or y.layout is torch.strided:
                        self.assertEqual(res.layout, torch.strided)
                    else:
                        self.assertEqual(res.layout, torch.sparse_csr)
                    self.assertEqual(res.to_dense(), expected)

                def _test(t, x, y):
                    _test_addmm(t, x, y)
                    _test_mm(x, y)

                if nnz0 is None:
                    nnz0 = random.randint(di * dk // 2, di * dk)
                t = torch.randn(di, dj, dtype=dtype, device=device)
                x = self.genSparseCSRTensor((di, dk), nnz0, device=device, dtype=dtype, index_dtype=index_dtype)
                y = torch.randn(dk, dj, dtype=dtype, device=device)
                _test(t, x, y)

                t = torch.randn(di, dj, dtype=dtype, device=device)
                x = self.genSparseCSCTensor((di, dk), nnz0, device=device, dtype=dtype, index_dtype=index_dtype)
                y = torch.randn(dk, dj, dtype=dtype, device=device)
                _test(t, x, y)

                if nnz1 is None:
                    nnz1 = random.randint(dk * dj // 2, dk * dj)
                t = torch.randn(di, dj, dtype=dtype, device=device)
                x = torch.randn(di, dk, dtype=dtype, device=device)
                y = self.genSparseCSRTensor((dk, dj), nnz1, device=device, dtype=dtype, index_dtype=index_dtype)
                _test(t, x, y)

                t = torch.randn(di, dj, dtype=dtype, device=device)
                x = torch.randn(di, dk, dtype=dtype, device=device)
                y = self.genSparseCSCTensor((dk, dj), nnz1, device=device, dtype=dtype, index_dtype=index_dtype)
                _test(t, x, y)

                x_shape, y_shape = x.shape, y.shape

                gen_csr_csc = [self.genSparseCSRTensor, self.genSparseCSCTensor]

                # Test mm({CSR, CSC}, {CSR, CSC})
                for gen_x, gen_y in itertools.product(gen_csr_csc, gen_csr_csc):
                    x = gen_x(x_shape, nnz0, device=device, dtype=dtype, index_dtype=index_dtype)
                    y = gen_y(y_shape, nnz1, device=device, dtype=dtype, index_dtype=index_dtype)
                    _test_mm(x, y)

        def test_empty_inputs(lhs_layout, rhs_layout):
            xd = torch.rand(10, 0, device=device, dtype=dtype)
            yd = xd.transpose(-2, -1)
            zd = torch.rand(0, 0, device=device, dtype=dtype)

            xls, yls, zls = (t.to_sparse(layout=lhs_layout) for t in (xd, yd, zd))
            xrs, yrs, zrs = (t.to_sparse(layout=rhs_layout) for t in (xd, yd, zd))

            for ls, rs, ld, rd in [(xls, yrs, xd, yd), (xls, zrs, xd, zd), (zls, yrs, zd, yd), (zls, zrs, zd, zd)]:
                res_sparse = ls @ rs
                res_dense = ld @ rd
                self.assertEqual(res_sparse.to_dense(), res_dense)

        def test_orthogonal_inputs(lhs_layout, rhs_layout):
            ones = torch.ones(2, 2, device=device, dtype=dtype)
            zeros = torch.zeros(2, 2, device=device, dtype=dtype)
            x = torch.cat((ones, zeros), -1).to_sparse(layout=lhs_layout)
            y = torch.cat((zeros, ones), -2).to_sparse(layout=rhs_layout)
            res = x @ y
            res_expected = torch.zeros(*res.shape, device=device, dtype=dtype, layout=res.layout)
            self.assertEqual(res, res_expected)

        for lhs_layout, rhs_layout in itertools.product([torch.sparse_csr, torch.sparse_csc], repeat=2):
            test_empty_inputs(lhs_layout, rhs_layout)
            test_orthogonal_inputs(lhs_layout, rhs_layout)

        for i in [2, 4]:
            for j in [2, 4, 7]:
                for k in [2, 3, 7]:
                    test_shape(i, j, k)
        test_shape(4, 4, 4, 0, 0)

    @skipCPUIfNoMklSparse
    @dtypes(*floating_and_complex_types())
    @dtypesIfCUDA(*floating_and_complex_types_and(
                  *[torch.half] if SM53OrLater and TEST_CUSPARSE_GENERIC else [],
                  *[torch.bfloat16] if SM80OrLater and TEST_CUSPARSE_GENERIC else []))
    @precisionOverride({torch.bfloat16: 1e-2, torch.float16: 1e-2})
    def test_sparse_mm(self, device, dtype):
        def test_shape(d1, d2, d3, nnz, transposed, index_dtype):
            if transposed:
                D = torch.randn(d3, d2, dtype=dtype, device=device).t_()
            else:
                D = torch.randn(d2, d3, dtype=dtype, device=device)
            S = self.genSparseCSRTensor((d1, d2), nnz, device=device, dtype=dtype, index_dtype=index_dtype)
            S_dense = S.to_dense()
            self.assertEqual(torch.sparse.mm(S, D), torch.mm(S_dense, D))

        for index_dtype in [torch.int32, torch.int64]:
            test_shape(7, 8, 9, 20, False, index_dtype)
            test_shape(7, 8, 9, 20, True, index_dtype)

    @dtypes(*floating_and_complex_types())
    @dtypesIfCUDA(*floating_and_complex_types_and(
                  *[torch.half] if SM53OrLater and TEST_CUSPARSE_GENERIC else [],
                  *[torch.bfloat16] if SM80OrLater and TEST_CUSPARSE_GENERIC else []))
    @precisionOverride({torch.bfloat16: 1e-2, torch.float16: 1e-2})
    def test_sparse_addmm(self, device, dtype):
        def test_shape(m, n, p, nnz, broadcast, index_dtype, alpha_beta=None):
            if alpha_beta is None:
                alpha = random.random()
                beta = random.random()
            else:
                alpha, beta = alpha_beta
            if broadcast:
                D1 = make_tensor((), dtype=dtype, device=device)
            else:
                D1 = make_tensor([n, p], dtype=dtype, device=device)
            D2 = make_tensor([m, p], dtype=dtype, device=device)
            S = self.genSparseCSRTensor([n, m], nnz, dtype=dtype, device=device, index_dtype=index_dtype)
            S_dense = S.to_dense()
            Y = torch.sparse.addmm(D1, S, D2, beta=beta, alpha=alpha)
            Y_dense = torch.addmm(D1, S_dense, D2, beta=beta, alpha=alpha)
            self.assertEqual(Y, Y_dense)

        for index_dtype in [torch.int32, torch.int64]:
            test_shape(7, 8, 9, 20, False, index_dtype, None)
            test_shape(7, 8, 9, 20, True, index_dtype, None)
            test_shape(7, 8, 9, 20, False, index_dtype, (1, 0))
            test_shape(7, 8, 9, 20, True, index_dtype, (1, 0))
            test_shape(7, 8, 9, 20, False, index_dtype, (1, 1))
            test_shape(7, 8, 9, 20, True, index_dtype, (1, 1))

    @skipCPUIfNoMklSparse
    @dtypes(*floating_and_complex_types())
    @precisionOverride({torch.double: 1e-8, torch.float: 1e-4, torch.bfloat16: 0.6,
                        torch.half: 1e-1, torch.cfloat: 1e-4, torch.cdouble: 1e-8})
    @dtypesIfCUDA(*floating_types_and(torch.complex64,
                                      *[torch.bfloat16] if SM80OrLater else [],
                                      *[torch.half] if SM53OrLater else [],
                                      *[torch.complex128] if CUSPARSE_SPMM_COMPLEX128_SUPPORTED else []))
    @sparse_compressed_nonblock_layouts()
    @skipCUDAIf(
        not _check_cusparse_spgemm_available(),
        "cuSparse Generic API SpGEMM is not available"
    )
    def test_addmm_all_sparse_csr(self, device, dtype, layout):
        M = torch.randn(10, 25, device=device).to(dtype)
        m1 = torch.randn(10, 50, device=device).to(dtype)
        m2 = torch.randn(50, 25, device=device).to(dtype)
        _test_addmm_addmv(self, torch.addmm, M, m1, m2, layout=layout, mode="all_sparse")

        # Test 0-strided
        M = torch.randn(10, 1, device=device).to(dtype).expand(10, 25)
        m1 = torch.randn(10, 1, device=device).to(dtype).expand(10, 50)
        m2 = torch.randn(50, 25, device=device).to(dtype)
        _test_addmm_addmv(self, torch.addmm, M, m1, m2, layout=layout, mode="all_sparse")

        # Test beta=0, M=nan
        M = torch.full((10, 25), float('nan'), device=device).to(dtype)
        m1 = torch.randn(10, 50, device=device).to(dtype)
        m2 = torch.randn(50, 25, device=device).to(dtype)
        _test_addmm_addmv(self, torch.addmm, M, m1, m2, beta=0, layout=layout, mode="all_sparse")

        # Test transpose
        for t1, t2, t3, t4 in itertools.product([True, False], repeat=4):
            def maybe_transpose(cond, m):
                if not cond:
                    return m
                return m.t().clone(memory_format=torch.contiguous_format).t()

            M = maybe_transpose(t1, torch.randn(10, 25, device=device).to(dtype))
            m1 = maybe_transpose(t2, torch.randn(10, 50, device=device).to(dtype))
            m2 = maybe_transpose(t3, torch.randn(50, 25, device=device).to(dtype))
            _test_addmm_addmv(self, torch.addmm, M, m1, m2, transpose_out=t4, layout=layout, mode="all_sparse")

    @onlyCPU
    @skipCPUIfNoMklSparse
    @dtypes(*floating_and_complex_types())
    @sparse_compressed_nonblock_layouts()
    def test_addmm_dense_result(self, device, dtype, layout):
        M = torch.randn(10, 25, device=device).to(dtype)
        m1 = torch.randn(10, 50, device=device).to(dtype)
        m2 = torch.randn(50, 25, device=device).to(dtype)
        _test_addmm_addmv(self, torch.addmm, M, m1, m2, layout=layout, mode="dense_result")

        # Test 0-strided
        M = torch.randn(10, 1, device=device).to(dtype).expand(10, 25)
        m1 = torch.randn(10, 1, device=device).to(dtype).expand(10, 50)
        m2 = torch.randn(50, 25, device=device).to(dtype)
        _test_addmm_addmv(self, torch.addmm, M, m1, m2, layout=layout, mode="dense_result")

        # Test beta=0, M=nan
        M = torch.full((10, 25), float('nan'), device=device).to(dtype)
        m1 = torch.randn(10, 50, device=device).to(dtype)
        m2 = torch.randn(50, 25, device=device).to(dtype)
        _test_addmm_addmv(self, torch.addmm, M, m1, m2, beta=0, layout=layout, mode="dense_result")

        # Test transpose
        for t1, t2, t3, t4 in itertools.product([True, False], repeat=4):
            def maybe_transpose(cond, m):
                if not cond:
                    return m
                return m.t().clone(memory_format=torch.contiguous_format).t()

            M = maybe_transpose(t1, torch.randn(10, 25, device=device).to(dtype))
            m1 = maybe_transpose(t2, torch.randn(10, 50, device=device).to(dtype))
            m2 = maybe_transpose(t3, torch.randn(50, 25, device=device).to(dtype))
            _test_addmm_addmv(self, torch.addmm, M, m1, m2, transpose_out=t4, layout=layout, mode="dense_result")

    @parametrize("k", [0, 1, 8])
    @parametrize("n", [0, 1, 10])
    @parametrize("m", [0, 1, 25])
    @skipCPUIfNoMklSparse
    @dtypes(*floating_and_complex_types())
    @dtypesIfCUDA(*floating_types_and(torch.complex64,
                                      *[torch.bfloat16] if SM80OrLater else [],
                                      *[torch.half] if SM53OrLater else [],
                                      *[torch.complex128]
                                      if CUSPARSE_SPMM_COMPLEX128_SUPPORTED or HIPSPARSE_SPMM_COMPLEX128_SUPPORTED
                                      else []))
    @precisionOverride({torch.double: 1e-8, torch.float: 1e-4, torch.bfloat16: 0.6,
                        torch.half: 1e-1, torch.cfloat: 1e-4, torch.cdouble: 1e-8})
    def test_addmm_sizes_all_sparse_csr(self, device, dtype, m, n, k):
        if (TEST_WITH_ROCM and k != 0 and n != 0 and m != 0):
            self.skipTest("Skipped on ROCm")
        M = torch.randn(n, m, device=device).to(dtype)
        m1 = torch.randn(n, k, device=device).to(dtype)
        m2 = torch.randn(k, m, device=device).to(dtype)
        _test_addmm_addmv(self, torch.addmm, M, m1, m2, layout=torch.sparse_csr, mode="all_sparse")

        M = torch.randn(n, m, device=device).to(dtype).to_sparse_csr()
        m1 = torch.randn(n, k + 1, device=device).to(dtype).to_sparse_csr()
        m2 = torch.randn(k, m, device=device).to(dtype).to_sparse_csr()
        self.assertRaisesRegex(RuntimeError, f"{n}x{k + 1}.*{k}x{m}", lambda: torch.addmm(M, m1, m2))
        self.assertRaisesRegex(RuntimeError, f"{n}x{k + 1}.*{k}x{m}", lambda: torch.mm(m1, m2))

    @skipCPUIfNoMklSparse
    @dtypes(torch.float)
    def test_addmm_errors(self, device, dtype):
        # test that the errors are the same for dense and sparse versions
        import re

        def test1(*, is_sparse):
            # shapes must be compatible for matrix multiplication
            a = make_tensor((2, 3), dtype=dtype, device=device)
            if is_sparse:
                a_sparse = a.to_sparse_csr()
                return torch.addmm(a, a_sparse, a)
            else:
                return torch.addmm(a, a, a)

        def test2(*, is_sparse):
            # mat2 must be a matrix
            a = make_tensor((2, 3), dtype=dtype, device=device)
            if is_sparse:
                a_sparse = a.to_sparse_csr()
                return torch.addmm(a, a_sparse, a.unsqueeze(0))
            else:
                return torch.addmm(a, a, a.unsqueeze(0))

        def test3(*, is_sparse):
            # the first input needs to be 1D or 2D
            a = make_tensor((3, 3), dtype=dtype, device=device)
            if is_sparse:
                a_sparse = a.to_sparse_csr()
                return torch.addmm(a.unsqueeze(0), a_sparse, a)
            else:
                return torch.addmm(a.unsqueeze(0), a, a)

        for test in (test1, test2, test3):
            try:
                test(is_sparse=False)
            except RuntimeError as msg:
                with self.assertRaisesRegex(RuntimeError, re.escape(str(msg))):
                    test(is_sparse=True)

    @skipCPUIfNoMklSparse
    @dtypes(torch.float)
    def test_mm_errors(self, device, dtype):
        # test that the errors are the same for dense and sparse versions
        import re

        def test1(*, is_sparse):
            # shapes must be compatible for matrix multiplication
            a = make_tensor((2, 3), dtype=dtype, device=device)
            if is_sparse:
                a_sparse = a.to_sparse_csr()
                return torch.mm(a_sparse, a)
            else:
                return torch.mm(a, a)

        def test2(*, is_sparse):
            # mat2 must be a matrix
            a = make_tensor((2, 3), dtype=dtype, device=device)
            if is_sparse:
                a_sparse = a.to_sparse_csr()
                return torch.mm(a_sparse, a.unsqueeze(0))
            else:
                return torch.mm(a, a.unsqueeze(0))

        for test in (test1, test2):
            try:
                test(is_sparse=False)
            except RuntimeError as msg:
                with self.assertRaisesRegex(RuntimeError, re.escape(str(msg))):
                    test(is_sparse=True)

    @sparse_compressed_nonblock_layouts()
    @dtypes(torch.float, torch.double)
    def test_add(self, device, layout, dtype):
        def _test_spadd_shape(nnz, shape):
            # sparse.to_dense() uses torch.add internally so if torch.add is wrong,
            # the dense tensor will be wrong but this test would still pass
            # there's a separate test that checks for the correctness of the .to_dense() call
            x = self.genSparseCompressedTensor(shape, nnz,
                                               dtype=dtype,
                                               device=device,
                                               index_dtype=torch.int32,
                                               layout=layout,
                                               blocksize=())
            y = torch.randn(*shape, dtype=dtype, device=device)
            r = random.random()

            res = torch.add(y, x, alpha=r)
            expected = y + r * x.to_dense()
            self.assertEqual(res, expected)
            res_perm = torch.add(x, y, alpha=r)
            self.assertEqual(res_perm, expected)

            # Non contiguous dense tensor
            s = list(shape)
            s[0] = shape[-1]
            s[-1] = shape[0]
            y = torch.randn(*s, dtype=torch.double, device=device)
            y.transpose_(0, len(s) - 1)
            r = random.random()

            res = torch.add(y, x, alpha=r)
            expected = y + r * x.to_dense()
            res_perm = torch.add(x, y, alpha=r)

            self.assertEqual(res, expected)
            self.assertEqual(res_perm, expected)


        ns = [2, 5]
        batch_shapes = [(), (2,), (2, 3)]
        for b, m, n in itertools.product(batch_shapes, ns, ns):
            _test_spadd_shape(0, (*b, m, n))
            _test_spadd_shape(m * n // 2, (*b, m, n))
            _test_spadd_shape(m * n, (*b, m, n))

    @dtypes(torch.float, torch.double)
    def test_mul(self, device, dtype):
        # TODO: This whole test should be migrated to OpInfos
        def _test_spadd_shape(fn, nnz, shape):
            x = self.genSparseCSRTensor(shape, nnz, dtype=dtype, device=device, index_dtype=torch.int32)
            y = self.genSparseCSRTensor(shape, nnz, dtype=dtype, device=device, index_dtype=torch.int32)

            # Forward comparison
            res_sparse_sparse = fn(y, x)
            res_dense_sparse = fn(y.to_dense(), x)
            res_sparse_dense = fn(y, x.to_dense())
            expected = fn(y.to_dense(), x.to_dense())
            self.assertEqual(res_sparse_sparse, expected)
            # TODO: While result of mul(dense, csr) is csr, it is not fully compressed.
            # That means it may contain materialized zeros, since the dense argument
            # is converted according to the sparsity pattern of csr. In the future
            # we might require the result to be fully compressed.
            self.assertEqual(res_dense_sparse, expected)
            self.assertEqual(res_sparse_dense, expected)

            # Grad comparison
            x = self.genSparseCSRTensor(shape, nnz, dtype=dtype, device=device, index_dtype=torch.int32)
            y = self.genSparseCSRTensor(shape, nnz, dtype=dtype, device=device, index_dtype=torch.int32)
            z = self.genSparseCSRTensor(shape, nnz, dtype=dtype, device=device, index_dtype=torch.int32)

            # csr * csr -> csr with csr, csr gradients
            x_a = x.clone().requires_grad_()
            y_a = y.clone().requires_grad_()

            fn(y_a, x_a).backward(z)

            x_dense_a = x.to_dense().requires_grad_()
            y_dense_a = y.to_dense().requires_grad_()

            fn(y_dense_a, x_dense_a).backward(z.to_dense())

            self.assertEqual(x_a.grad.layout, torch.sparse_csr)
            self.assertEqual(y_a.grad.layout, torch.sparse_csr)

            self.assertEqual(x_a.grad.to_dense(), x_dense_a.grad)
            self.assertEqual(y_a.grad.to_dense(), y_dense_a.grad)

            # TODO: Currently strided Tensors cannot have csr gradients
            # dense * csr -> csr with csr, dense gradients
            x_a = x.clone().requires_grad_()
            y_a = y.to_dense().clone().requires_grad_()
            err_msg = "Function MulBackward0 returned an invalid gradient at index 0 - expected layout Strided but got SparseCsr"
            with self.assertRaisesRegex(RuntimeError, err_msg):
                fn(y_a, x_a).backward(z)

            # csr * dense -> csr with dense, csr gradients
            x_a = x.to_dense().clone().requires_grad_()
            y_a = y.clone().requires_grad_()
            err_msg = "Function MulBackward0 returned an invalid gradient at index 1 - expected layout Strided but got SparseCsr"
            with self.assertRaisesRegex(RuntimeError, err_msg):
                fn(y_a, x_a).backward(z)

        _test_spadd_shape(torch.mul, 100, [100, 100])
        _test_spadd_shape(torch.mul, 0, [100, 100])
        _test_spadd_shape(torch.mul, 100, [100, 1])
        _test_spadd_shape(torch.mul, 100, [1, 100])

    # TODO: enable hybrid once to_dense supports it
    @parametrize('enable_hybrid', [False])
    @all_sparse_compressed_layouts()
    @dtypes(*all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half))
    def test_mul_scalar(self, layout, device, dtype, enable_hybrid):
        for sparse in self.generate_simple_inputs(
                layout, device=device, dtype=dtype, index_dtype=torch.int32, enable_hybrid=enable_hybrid):
            for scalar_dtype in all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half):
                # ComplexHalf is experimental
                if dtype is torch.half and scalar_dtype.is_complex:
                    continue

                scalar_t = torch.tensor(2, dtype=scalar_dtype)
                for scalar in (scalar_t, scalar_t.item()):
                    res_out = sparse.mul(scalar)
                    self.assertEqual(res_out, scalar * sparse)

                    res_dense_out = sparse.to_dense().mul(scalar)
                    # BUG: dispatcher ignores mul.Scalar(Tensor, Scalar)
                    # This issues is circumvented in the mul(Tensor, Tensor) kernel.
                    self.assertEqual(res_out, res_dense_out)

                    if dtype == torch.result_type(sparse, scalar):
                        res_in_dense = sparse.to_dense().mul_(scalar)
                        res_in = sparse.clone().mul_(scalar)
                        self.assertEqual(res_in, res_in_dense)
                        self.assertEqual(res_out, res_in)

    @skipCPUIfNoMklSparse
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    def test_sparse_add(self, device, dtype):
        def run_test(m, n, index_dtype):

            alpha = random.random()
            nnz1 = random.randint(0, m * n)
            nnz2 = random.randint(0, m * n)
            nnz3 = random.randint(0, m * n)

            if TEST_WITH_ROCM:
                # ROCm fails when nnz = 0
                nnz1, nnz2, nnz3 = max(1, nnz1), max(1, nnz2), max(1, nnz3)

            S1 = self.genSparseCSRTensor([m, n], nnz1, dtype=dtype, device=device, index_dtype=index_dtype)
            S2 = self.genSparseCSRTensor([m, n], nnz2, dtype=dtype, device=device, index_dtype=index_dtype)
            S3 = self.genSparseCSRTensor([m, n], nnz3, dtype=dtype, device=device, index_dtype=index_dtype)
            sparse_args = [S1, S2, S3]
            dense_args = [t.to_dense() for t in sparse_args]
            arg_idx = list(range(len(sparse_args)))
            out_idx = arg_idx + [None]

            for idx1, idx2, idx3 in itertools.product(arg_idx, arg_idx, out_idx):
                s1 = sparse_args[idx1]
                s2 = sparse_args[idx2]
                s3 = None if idx3 is None else sparse_args[idx3]
                d1 = dense_args[idx1]
                d2 = dense_args[idx2]
                d3 = None if idx3 is None else dense_args[idx3]

                expected = torch.add(d1, d2, alpha=alpha, out=d3)
                actual = torch.add(s1, s2, alpha=alpha, out=s3)
                self.assertEqual(actual.crow_indices().dtype, index_dtype)
                self.assertEqual(actual.col_indices().dtype, index_dtype)
                self.assertEqual(actual, expected)
                self.assertEqual(s3, d3)
                if s3 is not None:
                    self.assertEqual(s3.crow_indices().dtype, index_dtype)
                    self.assertEqual(s3.col_indices().dtype, index_dtype)

        for index_dtype in [torch.int32, torch.int64]:
            for m, n in itertools.product([3, 5], [3, 5]):
                run_test(m, n, index_dtype)

    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    def test_sparse_add_errors(self, device, dtype):
        def run_test(index_type):
            a = self.genSparseCSRTensor((2, 2), 3, dtype=dtype, device=device, index_dtype=index_dtype)
            b = self.genSparseCSRTensor((2, 1), 2, dtype=dtype, device=device, index_dtype=index_dtype)
            with self.assertRaisesRegex(RuntimeError, "Expected input tensors to have the same shape"):
                torch.add(a, b)

        for index_dtype in [torch.int32, torch.int64]:
            run_test(index_dtype)

    @skipCPUIfNoMklSparse
    @skipCUDAIf(
        not _check_cusparse_triangular_solve_available(),
        "cuSparse Generic API SpSV is not available"
    )
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
                        torch.float64: 1e-8, torch.complex128: 1e-8})
    def test_sparse_triangular_solve(self, device, dtype):

        def run_test(n, k, upper, unitriangular, transpose, zero):
            if not unitriangular:
                triangle_function = torch.triu if upper else torch.tril
            else:
                # Make sure diagonal elements are not materialized.
                # This is to exercise `unitriangular=True` not relying on
                # explicit presence of these indices.
                if upper:
                    def remove_diagonal(t):
                        return t.triu(-1)

                else:
                    def remove_diagonal(t):
                        return t.tril(-1)

                triangle_function = remove_diagonal

            make_A = torch.zeros if zero else make_tensor
            A = make_A((n, n), dtype=dtype, device=device)
            A = triangle_function(A)
            A_sparse = A.to_sparse_csr()
            B = make_tensor((n, k), dtype=dtype, device=device)

            expected = torch.triangular_solve(B, A, upper=upper, unitriangular=unitriangular, transpose=transpose)
            expected_X = expected.solution

            actual = torch.triangular_solve(B, A_sparse, upper=upper, unitriangular=unitriangular, transpose=transpose)
            actual_X = actual.solution
            actual_A_clone = actual.cloned_coefficient
            self.assertTrue(actual_A_clone.numel() == 0)
            if A_sparse._nnz() == 0:
                self.assertTrue(actual_X.isnan().all())
                return
            self.assertEqual(actual_X, expected_X)

            # test out with C contiguous strides
            out = torch.empty_strided((n, k), (k, 1), dtype=dtype, device=device)
            torch.triangular_solve(
                B, A_sparse,
                upper=upper, unitriangular=unitriangular, transpose=transpose, out=(out, actual_A_clone)
            )
            self.assertEqual(out, expected_X)

            # test out with F contiguous strides
            out = torch.empty_strided((n, k), (1, n), dtype=dtype, device=device)
            torch.triangular_solve(
                B, A_sparse,
                upper=upper, unitriangular=unitriangular, transpose=transpose, out=(out, actual_A_clone)
            )
            self.assertEqual(out, expected_X)
            self.assertEqual(out.stride(), (1, n))

            # test out with discontiguous strides
            out = torch.empty_strided((2 * n, k), (1, 2 * n), dtype=dtype, device=device)[::2]
            if n > 0 and k > 0:
                self.assertFalse(out.is_contiguous())
                self.assertFalse(out.t().is_contiguous())
            before_stride = out.stride()
            torch.triangular_solve(
                B, A_sparse,
                upper=upper, unitriangular=unitriangular, transpose=transpose, out=(out, actual_A_clone)
            )
            self.assertEqual(out, expected_X)
            self.assertEqual(out.stride(), before_stride)

        ks = [0, 1, 3]
        ns = [5, 3, 0]
        for (k, n), (upper, unitriangular, transpose, zero) in itertools.product(itertools.product(ks, ns),
                                                                                 itertools.product([True, False], repeat=4)):
            run_test(n, k, upper, unitriangular, transpose, zero)

    @skipCUDAIf(
        not _check_cusparse_sddmm_available(),
        "cuSparse Generic API SDDMM is not available"
    )
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
                        torch.float64: 1e-8, torch.complex128: 1e-8})
    def test_sampled_addmm(self, device, dtype):
        def run_test(c, a, b, op_a, op_b, *, alpha=None, beta=None):
            if dtype.is_complex:
                alpha = random.random() + 0.3j if alpha is None else alpha
                beta = random.random() + 0.6j if beta is None else beta
            else:
                alpha = random.random() if alpha is None else alpha
                beta = random.random() if beta is None else beta

            if op_a and a.shape == b.shape:
                a = a.mH
            if op_b and a.shape == b.shape:
                b = b.mH

            actual = torch.sparse.sampled_addmm(c, a, b, alpha=alpha, beta=beta)

            out = torch.sparse_csr_tensor(
                *map(torch.clone, (actual.crow_indices(), actual.col_indices())),
                torch.empty_like(actual.values()),
                size=actual.shape
            )
            torch.sparse.sampled_addmm(c, a, b, alpha=alpha, beta=beta, out=out)

            spy_c = torch.sparse_csr_tensor(c.crow_indices(), c.col_indices(), torch.ones_like(c.values()), size=c.shape)
            expected = alpha * (a @ b) * spy_c.to_dense() + beta * c.to_dense()
            self.assertEqual(actual.to_dense(), out.to_dense())
            self.assertEqual(actual.to_dense(), expected)

        mnk = list(itertools.product([2, 5], repeat=3))

        # Add a test case for size 0 a and b tensors
        mnk = mnk + [(5, 5, 0)]

        batch_shapes = [(), (2,), (2, 3)]
        tf = [True, False]
        for index_dtype in [torch.int32, torch.int64]:
            for (m, n, k), b, noncontiguous, bcast_c in itertools.product(mnk, batch_shapes, tf, tf):
                if bcast_c and len(b) == 0:
                    continue
                nnz = random.randint(0, m * n)
                c_batch = () if bcast_c else b
                c = self.genSparseCSRTensor((*c_batch, m, n), nnz, dtype=dtype, device=device, index_dtype=index_dtype)
                a = make_tensor((*b, m, k), dtype=dtype, device=device, noncontiguous=noncontiguous)
                b = make_tensor((*b, k, n), dtype=dtype, device=device, noncontiguous=noncontiguous)
                for op_a, op_b in itertools.product([True, False], repeat=2):
                    run_test(c, a, b, op_a, op_b)

    @skipCUDAIf(
        not _check_cusparse_sddmm_available(),
        "cuSparse Generic API SDDMM is not available"
    )
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    def test_sampled_addmm_autograd(self, device, dtype):
        from torch.testing._internal.common_methods_invocations import sample_inputs_sparse_sampled_addmm

        samples = list(sample_inputs_sparse_sampled_addmm(None, device, dtype, requires_grad=True))

        for sample, dense_covector in zip(samples, [True, False]):
            c = sample.input
            a = sample.args[0]
            b = sample.args[1]

            # Compute sparse result
            output = torch.sparse.sampled_addmm(c, a, b, **sample.kwargs)
            covector = torch.randn_like(output).to_dense() if dense_covector else torch.randn_like(output)
            output.backward(covector)

            # Compute dense result and compare with sparse result
            c1, a1, b1 = (x.detach().to_dense().requires_grad_(True) for x in [c, a, b])
            dense_output = sample.kwargs['alpha'] * (a1 @ b1) * torch.ones_like(c).to_dense() + sample.kwargs['beta'] * c1
            self.assertEqual(output, dense_output)
            dense_covector = covector.to_dense()
            dense_output.backward(dense_covector)
            self.assertEqual(c.grad, c1.grad)
            self.assertEqual(a.grad, a1.grad)
            self.assertEqual(b.grad, b1.grad)

    @onlyCUDA
    # It works on ROCm and CUDA issue is currently active
    @skipCUDAIf(not TEST_WITH_ROCM, "Causes CUDA memory exception, see https://github.com/pytorch/pytorch/issues/72177")
    @skipCUDAIf(
        not _check_cusparse_sddmm_available(),
        "cuSparse Generic API SDDMM is not available"
    )
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    @precisionOverride({torch.float32: 1e-3, torch.complex64: 1e-3,
                        torch.float64: 1e-8, torch.complex128: 1e-8})
    def test_sampled_addmm_zero_sized(self, device, dtype):
        def run_test(c, a, b):
            actual = torch.sparse.sampled_addmm(c, a, b)
            self.assertEqual(actual.shape, c.shape)

        for m, n, k in itertools.product([0, 5], repeat=3):
            c = torch.empty(m, n, dtype=dtype, device=device, layout=torch.sparse_csr)
            a = make_tensor((m, k), dtype=dtype, device=device)
            b = make_tensor((k, n), dtype=dtype, device=device)
            run_test(c, a, b)

    @onlyCUDA
    @skipCUDAIf(
        not _check_cusparse_sddmm_available(),
        "cuSparse Generic API SDDMM is not available"
    )
    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
    def test_sampled_addmm_errors(self, device, dtype):
        # test that the errors are the same for dense and sparse sampled versions
        # import re

        # shapes must be compatible for matrix multiplication
        a = make_tensor((2, 3), dtype=dtype, device=device)
        a_sparse = a.to_sparse_csr()
        with self.assertRaisesRegex(RuntimeError, r"cannot be multiplied"):
            torch.sparse.sampled_addmm(a_sparse, a, a)

        # mat1 must be a matrix
        with self.assertRaisesRegex(RuntimeError, r"Expected mat1 to be a matrix"):
            torch.sparse.sampled_addmm(a_sparse, a[..., 0, :], a)

        # mat2 must be a matrix
        with self.assertRaisesRegex(RuntimeError, r"Expected mat2 to be a matrix"):
            torch.sparse.sampled_addmm(a_sparse, a, a[..., 0, :])

        a = make_tensor((2, 2), dtype=dtype, device=device)
        b = make_tensor((3, 3), dtype=dtype, device=device)
        b_sparse = b.to_sparse_csr()
        with self.assertRaisesRegex(RuntimeError, r"self.shape\[-2\] must match mat1.shape\[-2\]"):
            torch.sparse.sampled_addmm(b_sparse, a, a)

        b = make_tensor((2, 3), dtype=dtype, device=device)
        b_sparse = b.to_sparse_csr()
        with self.assertRaisesRegex(RuntimeError, r"self.shape\[-1\] must match mat2.shape\[-1\]"):
            torch.sparse.sampled_addmm(b_sparse, a, a)

        a = make_tensor((2, 2), dtype=dtype, device=device)
        a_sparse = a.to_sparse_csr()
        with self.assertRaisesRegex(RuntimeError, r"Expected mat1 to have strided layout"):
            torch.sparse.sampled_addmm(a_sparse, a_sparse, a_sparse)

        with self.assertRaisesRegex(RuntimeError, r"Expected mat2 to have strided layout"):
            torch.sparse.sampled_addmm(a_sparse, a, a_sparse)

    @onlyCPU
    @dtypes(torch.float32, torch.float64, torch.bfloat16, torch.float16)
    @precisionOverride({torch.bfloat16: 0.01})
    def test_sparse_mm_reduce_sum(self, device, dtype):
        def run_test(m, n, k, nnz, train):
            sparse = self.genSparseCSRTensor((m, k), nnz, dtype=dtype, device=device, index_dtype=torch.int64)
            dense = sparse.to_dense()

            mat = torch.randn(k, n, dtype=dtype)
            ref_mat = mat.clone()

            if train:
                sparse.requires_grad_()
                mat.requires_grad_()
                dense.requires_grad_()
                ref_mat.requires_grad_()

            ref_out = torch.mm(dense, ref_mat)
            out = torch.sparse.mm(sparse, mat, 'sum')

            self.assertEqual(out, ref_out)

            if train:
                ref_out.sum().backward()
                out.sum().backward()

                grad_input = sparse.grad
                ref_grad_input = dense.grad
                grad_mat = mat.grad
                ref_grad_mat = ref_mat.grad

                self.assertEqual(grad_input.to_dense(), ref_grad_input)
                self.assertEqual(grad_mat, ref_grad_mat)

        run_test(4, 5, 4, 10, False)
        run_test(4, 4, 4, 16, True)

    @skipIfTorchDynamo()
    @onlyCPU
    @dtypes(torch.float32, torch.float64, torch.bfloat16, torch.float16)
    @precisionOverride({torch.bfloat16: 0.01, torch.float16: 0.01})
    def test_sparse_mm_reduce(self, device, dtype):
        def run_test(m, n, k, nnz, reduce_type, index_dtype, train):
            csr = self.genSparseCSRTensor((m, n), nnz, dtype=dtype, device=device, index_dtype=index_dtype)
            mat = torch.randn(n, k, dtype=dtype)
            ref_mat = mat.clone()
            ref_values = csr.values().clone()

            out_int32 = index_dtype == torch.int32
            coo_indices = torch._convert_indices_from_csr_to_coo(
                csr.crow_indices(),
                csr.col_indices(),
                out_int32=out_int32)
            row, col = coo_indices[0], coo_indices[1]

            def ref(row, col, val, mat):
                out = torch.zeros([m, k], dtype=dtype)
                weight = mat.index_select(0, col)
                src = weight.mul(val.view(-1, 1))
                index = row.view(-1, 1).expand_as(weight)
                index = index.to(dtype=torch.int64)
                # scatter_reduce expect index to be int64
                out.scatter_reduce_(0, index, src, reduce=reduce_type, include_self=False)
                return out

            if train:
                csr.requires_grad_()
                mat.requires_grad_()
                ref_values.requires_grad_()
                ref_mat.requires_grad_()

            ref_out = ref(row, col, ref_values, ref_mat)
            out = torch.sparse.mm(csr, mat, reduce_type)
            self.assertEqual(out, ref_out)

            if train and dtype not in (torch.bfloat16, torch.float16):
                ref_out.sum().backward()
                out.sum().backward()

                grad_values = csr.grad.values()
                grad_weight = mat.grad
                ref_grad_values = ref_values.grad
                ref_grad_weight = ref_mat.grad
                self.assertEqual(grad_values, ref_grad_values)
                self.assertEqual(grad_weight, ref_grad_weight)

        for train in [False, True]:
            for index_dtype in [torch.int32, torch.int64]:
                for reduce_type in ["sum", "mean", "amax", "amin"]:
                    # by setting nnz < M, create empty rows
                    run_test(3, 4, 11, 1, reduce_type, index_dtype, train)
                    run_test(3, 4, 11, 6, reduce_type, index_dtype, train)
                    run_test(3, 4, 11, 12, reduce_type, index_dtype, train)
                    # we are doing blocking with 4x vector length in the kernel,
                    # so need to test when K > 4x vector length
                    run_test(4, 7, 33, 13, reduce_type, index_dtype, train)

    @skipMeta
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_coo_csr_conversion(self, device, dtype):
        for m, n in itertools.product([5, 2, 0], [5, 2, 0]):
            size = (m, n)
            dense = make_tensor(size, dtype=dtype, device=device)
            coo_sparse = dense.to_sparse()
            csr_sparse = coo_sparse.to_sparse_csr()

            self.assertEqual(csr_sparse.to_dense(), dense)

    @skipMeta
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_csr_coo_conversion(self, device, dtype):
        for m, n in itertools.product([5, 2, 0], [5, 2, 0]):
            size = (m, n)
            dense = make_tensor(size, dtype=dtype, device=device)
            csr_sparse = dense.to_sparse_csr()
            coo_sparse = csr_sparse.to_sparse()

            self.assertEqual(coo_sparse.to_dense(), dense)

    # Currently, there is no rule in PyTorch for filling zeros in the outputs
    #   from operations on Sparse CSR tensors. Hence only those operators are supported
    #   which have 0->0 correspondence, example: sin(0) = 0, tan(0) = 0 but
    #   cos(0) = 1 (and hence it's not supported).
    # Note: here, we do this test only for unary operators
    @ops(sparse_csr_unary_ufuncs)
    def test_zero_to_zero_correspondence_unary(self, device, dtype, op):
        zero = torch.zeros((1, 2), dtype=dtype, device=device)
        tensor_explicit_zeros = torch.sparse_csr_tensor([0, 1], [1], [0], dtype=dtype, device=device)

        output_zero = op(zero)
        expected_zero = zero.to(output_zero.dtype)

        output_explicit_zeros = op(tensor_explicit_zeros).to_dense()
        expected_explicit_zeros = tensor_explicit_zeros.to_dense().to(output_explicit_zeros.dtype)

        for (output, expected) in [
                (output_zero, expected_zero),
                (output_explicit_zeros, expected_explicit_zeros)
        ]:
            self.assertEqual(output, expected, f"This operator ({op.name}) should not be supported for "
                             "Sparse CSR as it breaks 0->0 correspondence.")

        for inp in [zero.to_sparse_csr(), tensor_explicit_zeros]:
            self.assertEqual(op(inp).values().numel(), inp.values().numel(),
                             f"{op.name} fails to preserve sparsity pattern.")

    @ops(sparse_csr_unary_ufuncs)
    def test_sparse_csr_unary_out(self, device, dtype, op):
        samples = op.sample_inputs(device, dtype)

        if not op.supports_out:
            self.skipTest("Skipped! Out not supported")

        for sample in samples:
            assert torch.is_tensor(sample.input)
            # Sparse CSR only supports 2D tensors as inputs
            # Fail early to prevent silent success with this test
            if sample.input.ndim != 2:
                raise ValueError("Expected 2D tensor but got tensor with dimension: {sample.input.ndim}.")

            sample.input = sample.input.to_sparse_csr()
            expect = op(sample.input, *sample.args, **sample.kwargs)

            out = self.genSparseCSRTensor(sample.input.size(), sample.input._nnz(),
                                          device=sample.input.device, dtype=expect.dtype,
                                          index_dtype=sample.input.crow_indices().dtype)
            op(sample.input, *sample.args, **sample.kwargs, out=out)

            self.assertEqual(out, expect)

    @ops(sparse_csr_unary_ufuncs)
    def test_sparse_csr_unary_inplace(self, device, dtype, op):
        samples = op.sample_inputs(device, dtype)

        if op.inplace_variant is None:
            self.skipTest("Skipped! Inplace variant not supported!")

        for sample in samples:
            assert torch.is_tensor(sample.input)
            # Sparse CSR only supports 2D tensors as inputs
            # Fail early to prevent silent success with this test
            if sample.input.ndim != 2:
                raise ValueError("Expected 2D tensor but got tensor with dimension: {sample.input.ndim}.")

            sample.input = sample.input.to_sparse_csr()
            expect = op(sample.input, *sample.args, **sample.kwargs)

            if not torch.can_cast(expect.dtype, dtype):
                with self.assertRaisesRegex(RuntimeError, "result type"):
                    op.inplace_variant(sample.input, *sample.args, **sample.kwargs)
                continue

            if sample.input.is_complex() and op.name == "abs":
                with self.assertRaisesRegex(RuntimeError, "not supported"):
                    op.inplace_variant(sample.input, *sample.args, **sample.kwargs)
                continue

            actual = op.inplace_variant(sample.input, *sample.args, **sample.kwargs)

            self.assertIs(actual, sample.input)
            self.assertEqual(actual, expect)

    @skipIfTorchDynamo("Not a TorchDynamo suitable test")
    @ops(sparse_csr_unary_ufuncs, dtypes=OpDTypes.supported, allowed_dtypes=[torch.double, torch.cdouble])
    def test_autograd_sparse_csr_unary(self, device, dtype, op):
        if op.name not in UNARY_EWISE_CSR_ALLOW_AUTOGRAD:
            self.skipTest(f"Skipped! Unary op {op.name} not supported with CSR input and autograd")

        samples = list(op.sample_inputs(device, dtype))

        # Fail early to prevent silent success with this test
        ndims_equals_2d = (s.input.ndim == 2 for s in samples)
        if not any(ndims_equals_2d):
            raise ValueError("Expected at least one 2D tensor in samples.")

        for sample in samples:
            # We must skip samples of low dimensionality, we can't covert them to sparsed compressed layouts
            if sample.input.ndim < 2:
                continue
            sparse_input = sample.input.to_sparse_csr().requires_grad_(True)

            def fn(input):
                output = op.gradcheck_wrapper(op.get_op(), input, *sample.args, **sample.kwargs)
                if sample.output_process_fn_grad is not None:
                    return sample.output_process_fn_grad(output)
                return output

            # Compute sparse result
            output = fn(sparse_input)
            covector = torch.randn_like(output)
            output.backward(covector)
            self.assertTrue(torch.is_tensor(sparse_input.grad))
            self.assertTrue(sparse_input.grad.is_sparse_csr)

            # Compute dense result and compare with sparse result
            dense_input = sparse_input.detach().to_dense().requires_grad_(True)
            dense_output = fn(dense_input)
            dense_covector = covector.to_dense()
            dense_output.backward(dense_covector)
            self.assertEqual(sparse_input.grad, dense_input.grad)

    @skipCUDAIf(
        not _check_cusparse_sddmm_available(),
        "cuSparse Generic API SDDMM is not available"
    )
    @dtypes(torch.float64)
    def test_autograd_dense_output_addmm(self, device, dtype):
        from torch.testing._internal.common_methods_invocations import sample_inputs_addmm

        samples = list(sample_inputs_addmm(None, device, dtype, requires_grad=True))

        # Fail early to prevent silent success with this test
        ndims_equals_2d = (s.args[0].ndim == 2 for s in samples)
        if not any(ndims_equals_2d):
            raise ValueError("Expected at least one 2D tensor in samples to convert to sparse.")

        for sample in samples:
            a = sample.args[0].relu().to_sparse_csr()
            if sample.args[0].shape == sample.args[1].shape:
                import warnings
                warnings.warn("Broken for square matrices, see https://github.com/pytorch/pytorch/issues/116565")
                continue

            # This path tests the autograd path wrt dense inputs
            for addmm in [torch.addmm, torch.sparse.addmm]:

                def fn(c, b):
                    output = addmm(c, a, b, **sample.kwargs)
                    if sample.output_process_fn_grad is not None:
                        return sample.output_process_fn_grad(output)
                    return output

                self.assertTrue(torch.autograd.gradcheck(fn, [sample.input, sample.args[1]], fast_mode=True))

                # noncontiguous
                c = make_tensor(sample.input.shape, device=device, dtype=dtype, noncontiguous=True, requires_grad=True)
                b = make_tensor(sample.args[1].shape, device=device, dtype=dtype, noncontiguous=True, requires_grad=True)
                self.assertTrue(torch.autograd.gradcheck(fn, [c, b], fast_mode=True))

                # Now test the autograd path wrt sparse inputs
                for reverse in [True, False]:
                    c, b = sample.input, sample.args[1]
                    if reverse and a.shape != b.shape:
                        continue

                    def fn(a):
                        inputs = (c, b, a) if reverse else (c, a, b)
                        output = addmm(*inputs, **sample.kwargs)
                        if sample.output_process_fn_grad is not None:
                            return sample.output_process_fn_grad(output)
                        return output

                    # gradcheck doesn't work for sparse CSR yet, compare against dense path
                    # Compute sparse result
                    a = a.detach().requires_grad_(True)
                    output = fn(a)
                    covector = torch.randn_like(output)
                    output.backward(covector)
                    self.assertTrue(torch.is_tensor(a.grad))
                    if addmm == torch.sparse.addmm:
                        self.assertTrue(a.grad.is_sparse_csr)
                    else:
                        self.assertTrue(a.grad.layout == torch.strided)

                    # Compute dense result and compare with sparse result
                    dense_a = a.detach().to_dense().requires_grad_(True)
                    dense_output = fn(dense_a)
                    self.assertEqual(output, dense_output)
                    dense_covector = covector.to_dense()
                    dense_output.backward(dense_covector)

                    if addmm == torch.sparse.addmm:
                        self.assertEqual(a.grad, dense_a.grad.sparse_mask(a))
                    else:
                        self.assertEqual(a.grad, dense_a.grad)

    @skipCPUIfNoMklSparse
    @dtypes(torch.float64)
    def test_autograd_dense_output_addmv(self, device, dtype):
        from torch.testing._internal.common_methods_invocations import sample_inputs_addmv

        samples = list(sample_inputs_addmv(None, device, dtype, requires_grad=True))

        # Fail early to prevent silent success with this test
        ndims_equals_2d = (s.args[0].ndim == 2 for s in samples)
        if not any(ndims_equals_2d):
            raise ValueError("Expected at least one 2D tensor in samples to convert to sparse.")

        for sample in samples:
            # TODO: Remove detach once we have autograd support for CSR input
            a = sample.args[0].to_sparse_csr().detach()

            def fn(c, b):
                output = torch.addmv(c, a, b, **sample.kwargs)
                if sample.output_process_fn_grad is not None:
                    return sample.output_process_fn_grad(output)
                return output

            self.assertTrue(torch.autograd.gradcheck(fn, [sample.input, sample.args[1]], fast_mode=True))

            # noncontiguous
            c = make_tensor(sample.input.shape, device=device, dtype=dtype, noncontiguous=True, requires_grad=True)
            b = make_tensor(sample.args[1].shape, device=device, dtype=dtype, noncontiguous=True, requires_grad=True)
            self.assertTrue(torch.autograd.gradcheck(fn, [c, b], fast_mode=True))

    @skipIfTorchDynamo("Not a TorchDynamo suitable test")
    @ops(binary_ops_with_dense_output, dtypes=OpDTypes.supported, allowed_dtypes=[torch.double, ])
    def test_autograd_dense_output(self, device, dtype, op):
        if op.name == "mv" and no_mkl_sparse and self.device_type == 'cpu':
            self.skipTest("MKL Sparse is not available")

        samples = list(op.sample_inputs(device, dtype, requires_grad=True))

        # Fail early to prevent silent success with this test
        ndims_equals_2d = (s.input.ndim == 2 for s in samples)
        if not any(ndims_equals_2d):
            raise ValueError("Expected at least one 2D tensor in samples.")

        # Here we assume that the signature is op(sparse_input, dense_input) -> dense_output
        for sample in samples:
            # TODO: Remove detach once we have autograd support for CSR input
            sparse_input = sample.input.to_sparse_csr().detach()

            def fn(*args):
                output = op.gradcheck_wrapper(op.get_op(), sparse_input, *args, **sample.kwargs)
                if sample.output_process_fn_grad is not None:
                    return sample.output_process_fn_grad(output)
                return output

            self.assertTrue(torch.autograd.gradcheck(fn, sample.args, fast_mode=True))

            # noncontiguous
            args = [make_tensor(a.shape, device=device, dtype=dtype, noncontiguous=True, requires_grad=True) for a in sample.args]
            self.assertTrue(torch.autograd.gradcheck(fn, args, fast_mode=True))

    @dtypes(*all_types_and_complex())
    def test_direct_coo_csr_conversion(self, device, dtype):
        for m, n in itertools.product([5, 2, 0], [5, 2, 0]):
            size = (m, n)
            dense = make_tensor(size, dtype=dtype, device=device)
            coo_sparse = dense.to_sparse_coo()

            self.assertEqual(coo_sparse.to_sparse_csr().to_sparse_coo(), coo_sparse)

    @skipMeta
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_sum(self, device, dtype):
        def run_test(shape, nnz, index_type):
            a = self.genSparseCSRTensor(shape, nnz, dtype=dtype, device=device, index_dtype=index_dtype)
            self.assertEqual(a.sum(), a.values().sum())
            if dtype in floating_types():
                a.requires_grad_(True)
                a.sum().backward()
                self.assertEqual(a.grad, torch.ones(shape, dtype=dtype, device=device))
        for shape, index_dtype in itertools.product(
                [(10, 5), (10, 10)],
                [torch.int32, torch.int64]):
            run_test(shape, 0, index_dtype)
            run_test(shape, max(shape), index_dtype)
            run_test(shape, shape[0] * shape[1], index_dtype)

    @skipIfTorchDynamo()
    @skipMeta
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    @all_sparse_compressed_layouts()
    def test_transpose(self, device, dtype, layout):

        def _check_transpose_view(subject, transpose):
            self.assertTrue(transpose.values()._is_view())
            self.assertTrue(transpose._is_view())
            self.assertTrue(transpose._base is subject)

        def _check_layout_invariants(transpose):
            self.assertEqual(transpose.device, torch.device(device))
            compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[transpose.layout]
            compressed_indices, plain_indices = compressed_indices_mth(transpose), plain_indices_mth(transpose)
            torch._validate_sparse_compressed_tensor_args(compressed_indices, plain_indices, transpose.values(),
                                                          transpose.shape, transpose.layout)

        def check_good_transpose(subject, subject_dense, dim0, dim1, expected_layout):
            transpose = subject.transpose(dim0, dim1)
            # correct layout
            self.assertEqual(transpose.layout, expected_layout)
            # transpose must be return a view
            _check_transpose_view(subject, transpose)
            # result uses unsafe construction, so we check invariants
            _check_layout_invariants(transpose)
            self.assertEqual(transpose.to_dense(), subject_dense.transpose(dim0, dim1))

            round_trip = transpose.transpose(dim0, dim1)
            self.assertEqual(round_trip.layout, subject.layout)
            # transpose must be return a view
            _check_transpose_view(subject, round_trip)
            # result uses unsafe construction, so we check invariants
            _check_layout_invariants(round_trip)
            self.assertEqual(round_trip.to_dense(), subject_dense)

        def check_same_dim_transpose(subject, subject_dense, dim):
            transpose = subject.transpose(dim, dim)
            # correct layout
            self.assertEqual(transpose.layout, subject.layout)
            # transpose must be return a view
            _check_transpose_view(subject, transpose)
            # result uses unsafe construction, so we check invariants
            _check_layout_invariants(transpose)
            self.assertEqual(transpose.to_dense(), subject_dense)

        def check_dim_type_mismatch_throws(subject, name0, dim0, name1, dim1):
            mismatch_name = f"{dim0}\\({name0}\\) and {dim1}\\({name1}\\)"
            err = r"transpose\(\): can only transpose dimensions of the same type \(Batch, Sparse, Dense\), got " + mismatch_name

            with self.assertRaisesRegex(RuntimeError, err):
                subject.transpose(dim0, dim1)

        def run_test(shape, nnz, index_type, n_dense, blocksize=()):
            subject = self.genSparseCompressedTensor(shape,
                                                     nnz,
                                                     layout=layout,
                                                     device=device,
                                                     index_dtype=index_type,
                                                     blocksize=blocksize,
                                                     dense_dims=n_dense,
                                                     dtype=dtype)


            sparse0 = len(shape) - n_dense - 1
            sparse1 = sparse0 - 1

            dense0 = sparse0 + 1 if n_dense > 0 else None
            dense1 = dense0 + 1 if n_dense > 1 else None

            n_batch = len(shape) - n_dense - 2
            batch0 = sparse1 - 1 if n_batch > 0 else None
            batch1 = 0 if n_batch > 1 else None

            sparse_dims = (sparse0, sparse1)
            dense_dims = (dense0, dense1)
            batch_dims = (batch0, batch1)

            named0 = [(name, d[0]) for name, d in zip(["Batch", "Sparse", "Dense"], (batch_dims, sparse_dims, dense_dims))]
            named1 = [(name, d[1]) for name, d in zip(["Batch", "Sparse", "Dense"], (batch_dims, sparse_dims, dense_dims))]

            flipped_layout = {
                torch.sparse_csr: torch.sparse_csc,
                torch.sparse_csc: torch.sparse_csr,
                torch.sparse_bsr: torch.sparse_bsc,
                torch.sparse_bsc: torch.sparse_bsr
            }[layout]
            if n_dense > 0:
                # expect all transpose to throw
                for (name0, dim0), (name1, dim1) in itertools.product(named0, named1):
                    msg = r"transpose\(\): hybrid sparse compressed tensors with dense dimensions are not supported"
                    if (dim0 is not None) and (dim1 is not None):
                        with self.assertRaisesRegex(RuntimeError, msg):
                            subject.transpose(dim0, dim1)
            else:
                subject_dense = subject.to_dense()
                for (name0, dim0), (name1, dim1) in itertools.product(named0, named1):
                    if dim0 is not None:
                        check_same_dim_transpose(subject, subject_dense, dim0)

                        if dim1 is not None:
                            if name0 == name1:
                                expected_layout = flipped_layout if name0 == "Sparse" else layout
                                check_good_transpose(subject, subject_dense, dim0, dim1, expected_layout)
                            else:
                                check_dim_type_mismatch_throws(subject, name0, dim0, name1, dim1)

        # batch/sparse, sparse/dense only and full hybrid cases
        shape_ndense = list(itertools.product([(2, 4, 6, 2), (10, 6, 4, 2), (2, 4, 4, 2, 6)], [0, 1, 2]))
        # sparse only cases
        shape_ndense += [[(4, 8), 0], [(2, 2), 0], [(8, 4), 0]]
        for (shape, n_dense), index_dtype in itertools.product(shape_ndense, [torch.int32, torch.int64]):
            n_batch = len(shape) - n_dense - 2
            sparse_shape = shape[n_batch: n_batch + 2]
            if layout in (torch.sparse_bsr, torch.sparse_bsc):
                # for blocked all combinations of 2,1 should be valid blocksizes
                run_test(shape, 0, index_dtype, n_dense, blocksize=(2, 2))
                run_test(shape, max(sparse_shape), index_dtype, n_dense, blocksize=(2, 2))
                run_test(shape, sparse_shape[0] * sparse_shape[1], index_dtype, n_dense, blocksize=(2, 2))
                # repeat the realistic sparseity case with varried block sizes
                run_test(shape, max(sparse_shape), index_dtype, n_dense, blocksize=(2, 1))
                run_test(shape, max(sparse_shape), index_dtype, n_dense, blocksize=(1, 2))
                run_test(shape, max(sparse_shape), index_dtype, n_dense, blocksize=(1, 1))
            else:
                run_test(shape, 0, index_dtype, n_dense)
                run_test(shape, max(sparse_shape), index_dtype, n_dense)
                run_test(shape, sparse_shape[0] * sparse_shape[1], index_dtype, n_dense)

    # TODO: This is a stopgap for a rigorous extension of our autograd tests
    # to test the functionality of detach
    @skipMeta
    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
    def test_exercise_detach(self, device, dtype):
        shape = (3, 3)
        nnz = 4
        for index_dtype in [torch.int32, torch.int64]:
            inp = self.genSparseCSRTensor(shape, nnz, dtype=dtype, device=device, index_dtype=index_dtype)
            detached_inp = inp.detach()
            self.assertEqual(inp, detached_inp)

    def _construct_sp_matrix(self, tensor, layout, blocksize=(2, 2)):
        if tensor.layout in [torch.sparse_coo, torch.sparse_csr, torch.sparse_csc, torch.strided]:
            tensor = tensor.to_dense()
        else:
            raise NotImplementedError(repr(tensor))
        if layout is torch.sparse_csr:
            return sp.csr_matrix(tensor.cpu().numpy())
        if layout is torch.sparse_csc:
            return sp.csc_matrix(tensor.cpu().numpy())
        if layout is torch.sparse_bsr:
            return sp.bsr_matrix(tensor.cpu().numpy(), blocksize=blocksize).sorted_indices()
        if layout is torch.sparse_bsc:
            # SciPy doesn't have native BSC support - but our tests don't need the full
            # functionality so fake it by using a transposed BSR matrix.
            class FakeBscMatrix:
                def __init__(self, matrix):
                    self._matrix = matrix
                    self.shape = tuple(reversed(matrix.shape))
                    self.indptr = matrix.indptr
                    self.indices = matrix.indices
                    self.data = [x.transpose() for x in matrix.data]

                @staticmethod
                def from_matrix(matrix, blocksize):
                    blocksize = tuple(reversed(blocksize))
                    matrix = matrix.transpose()
                    return FakeBscMatrix(sp.bsr_matrix(matrix, blocksize=blocksize))

                def sorted_indices(self):
                    sub = self._matrix.sorted_indices()
                    return FakeBscMatrix(sub)

            return FakeBscMatrix.from_matrix(tensor.cpu().numpy(), blocksize=blocksize).sorted_indices()
        raise NotImplementedError(repr(tensor))

    @skipMeta
    @all_sparse_compressed_layouts('to_layout')
    @all_sparse_compressed_layouts('from_layout')
    def test_compressed_layout_conversions_coverage(self, device, from_layout, to_layout):
        """This test performs a smoke test for covered conversion and verifies
        that an exception is thrown for unsupported conversions.

        TODO: This test covers a subset of
        TestSparseAny.test_to_sparse tests and can be
        eliminated. Keeping the test until the new
        `Tensor.to_sparse(*, layout, blocksize)` has landed.
        """

        allowed_pairwise_layouts_sets = {
            frozenset({torch.sparse_csc}),
            frozenset({torch.sparse_csr}),
            frozenset({torch.sparse_csc, torch.sparse_csr}),
            frozenset({torch.sparse_csc, torch.sparse_bsc}),
            frozenset({torch.sparse_csc, torch.sparse_bsr}),
            frozenset({torch.sparse_csr, torch.sparse_bsc}),
            frozenset({torch.sparse_csr, torch.sparse_bsr}),
            frozenset({torch.sparse_bsc}),
            frozenset({torch.sparse_bsr}),
            frozenset({torch.sparse_bsc, torch.sparse_bsr}),
        }
        block_layouts = (torch.sparse_bsr, torch.sparse_bsc)

        def _to_from_layout(layout_a, layout_b, a):
            expect_error = True
            if {layout_a, layout_b} in allowed_pairwise_layouts_sets:
                expect_error = False

            # BSR -> CSR is not yet supported
            if (layout_a, layout_b) == (torch.sparse_bsr, torch.sparse_csr):
                expect_error = True
            # BSR -> CSC is not yet supported
            if (layout_a, layout_b) == (torch.sparse_bsr, torch.sparse_csc):
                expect_error = True
            # BSC -> CSR is not yet supported
            if (layout_a, layout_b) == (torch.sparse_bsc, torch.sparse_csr):
                expect_error = True
            # BSC -> CSC is not yet supported
            if (layout_a, layout_b) == (torch.sparse_bsc, torch.sparse_csc):
                expect_error = True
            # CSR -> BSR only works for non-batched inputs
            if (layout_a, layout_b) == (torch.sparse_csr, torch.sparse_bsr):
                if a.dim() > 2:
                    expect_error = True
            # CSR -> BSC only works for non-batched inputs
            if (layout_a, layout_b) == (torch.sparse_csr, torch.sparse_bsc):
                if a.dim() > 2:
                    expect_error = True
            # CSC -> BSR only works for non-batched inputs
            if (layout_a, layout_b) == (torch.sparse_csc, torch.sparse_bsr):
                if a.dim() > 2:
                    expect_error = True
            # CSC -> BSC only works for non-batched inputs
            if (layout_a, layout_b) == (torch.sparse_csc, torch.sparse_bsc):
                if a.dim() > 2:
                    expect_error = True

            blocksize_a = (1, 1) if layout_a in {torch.sparse_bsr, torch.sparse_bsc} else None
            blocksize_b = (1, 1) if layout_b in {torch.sparse_bsr, torch.sparse_bsc} else None
            b = a.to_sparse(layout=layout_a, blocksize=blocksize_a)
            if expect_error:
                with self.assertRaises(RuntimeError):
                    b.to_sparse(layout=layout_b, blocksize=blocksize_b)
            else:
                c = b.to_sparse(layout=layout_b, blocksize=blocksize_b)
                self.assertEqual(a.to_dense(), c.to_dense())

                # change of blocksize upon conversion is not yet supported.
                if b.layout in block_layouts:
                    for block_layout in block_layouts:
                        with self.assertRaisesRegex(RuntimeError,
                                                    "conversion from.*to.*with blocksize changed from.*to.*is not supported"):
                            b.to_sparse(layout=block_layout, blocksize=(3, 3))

        batch_dims = [(), (2,), (2, 2), (2, 2, 2)]
        sparse_dims = (6, 12)
        for batch_dim in batch_dims:
            a = make_tensor(batch_dim + sparse_dims, dtype=torch.float, device=device)
            _to_from_layout(from_layout, to_layout, a)

    @skipMeta
    @all_sparse_compressed_layouts()
    @batched_nonbatched()
    @hybrid_nonhybrid()
    @unittest.skipIf(not TEST_SCIPY, "SciPy not found")
    def test_dense_to_from_sparse_compressed(self, device, hybrid, batched, layout):
        """This test tests conversion from dense to/from CSR and CSC
        by comparing to SciPy's implementation.

        Here we test only those conversion combinations that SciPy
        supports to ensure that PyTorch conversions are in the same
        page with SciPy.  Independent from SciPy, all conversion
        combinations are tested in TestSparseAny.test_to_sparse.
        """

        blocked_layouts = (torch.sparse_bsr, torch.sparse_bsc)

        # helpers

        def _check_against_scipy_matrix(pt_matrix, dense, blocksize, **kwargs):
            # scipy has no bsc layout, so we check against the bsr layout of the tranposed dense
            if layout == torch.sparse_bsc:
                sp_matrix = self._construct_sp_matrix(dense.t(), layout=torch.sparse_bsr, blocksize=blocksize[::-1])
            else:
                sp_matrix = self._construct_sp_matrix(dense, layout=layout, blocksize=blocksize)

            compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[layout]

            self.assertEqual(layout, pt_matrix.layout)
            if layout == torch.sparse_bsc:
                self.assertEqual(sp_matrix.shape[::-1], pt_matrix.shape)
            else:
                self.assertEqual(sp_matrix.shape, pt_matrix.shape)

            self.assertEqual(torch.tensor(sp_matrix.indptr, dtype=torch.int64), compressed_indices_mth(pt_matrix))
            self.assertEqual(torch.tensor(sp_matrix.indices, dtype=torch.int64), plain_indices_mth(pt_matrix))
            if layout == torch.sparse_bsc:
                # we must tranpose the blocks before comparing
                self.assertEqual(torch.tensor(sp_matrix.data), pt_matrix.values().transpose(-2, -1))
            else:
                self.assertEqual(torch.tensor(sp_matrix.data), pt_matrix.values())

        def _check_hybrid_matrix(pt_matrix, dense, blocksize, **kwargs):
            # Calculate COO indices for sparse matrix.
            compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[layout]
            compressed_indices = compressed_indices_mth(pt_matrix)
            plain_indices = plain_indices_mth(pt_matrix)
            coo_indices = torch._convert_indices_from_csr_to_coo(compressed_indices, plain_indices)
            row_indices, col_indices = {
                torch.sparse_csr: (coo_indices[0, ], coo_indices[1, ]),
                torch.sparse_csc: (coo_indices[1, ], coo_indices[0, ]),
                torch.sparse_bsr: (coo_indices[0, ], coo_indices[1, ]),
                torch.sparse_bsc: (coo_indices[1, ], coo_indices[0, ]),
            }[pt_matrix.layout]

            # If sparse matrix layout blocked, rearrange dense matrix
            # so that the shape past first two dimensions match the
            # shape of sparse matrix values.
            dense_to_check = dense
            if blocksize:
                dense_shape = dense.shape
                dense_to_check_shape = (dense.shape[0] // blocksize[0],
                                        blocksize[0],
                                        dense.shape[1] // blocksize[1],
                                        blocksize[1]) + dense.shape[2:]
                dense_to_check = dense_to_check.reshape(dense_to_check_shape).transpose(1, 2)

            # Verify that non-zero values of the sparse matrix are
            # equal to corresponding values of the dense matrix.
            self.assertEqual(pt_matrix.values(), dense_to_check[row_indices, col_indices])

            # Verify that the remaining elements of the dense matrix
            # are 0, i.e. that dense are sparse matrix are fully
            # equal.
            mask = torch.ones_like(dense_to_check, dtype=torch.bool)
            mask[row_indices, col_indices] = False
            self.assertTrue(torch.all(torch.masked_select(dense_to_check, mask) == 0))

        def _check_batched(pt_tensor, dense, check_batch=None, batch_shape=(), blocksize=(), **kwargs):
            self.assertEqual(layout, pt_tensor.layout)
            self.assertEqual(pt_tensor.shape, dense.shape)
            compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[layout]
            for batch_index in np.ndindex(batch_shape):
                pt_matrix = pt_tensor[batch_index]
                dense_matrix = dense[batch_index]
                dense_dim = pt_matrix.dim() - 2
                dense_matrix_pt = dense_matrix.to_sparse(layout=layout,
                                                         blocksize=blocksize or None,
                                                         dense_dim=dense_dim)
                # sanity check, selecting batch of to_<layout> and dense[batch].to_<layout> should give the same result
                self.assertEqual(pt_matrix, dense_matrix_pt)
                check_batch(pt_matrix, dense_matrix, blocksize, **kwargs)

        def _generate_subject(sparse_shape, batch_shape, hybrid_shape):
            shape = batch_shape + sparse_shape + hybrid_shape
            n_batch_dim = len(batch_shape)
            n_hybrid_dim = len(hybrid_shape)
            # generate a dense tensor
            dense = make_tensor(shape, dtype=torch.float, device=device)

            # introduce some sparsty, mask is sparse shape, element applies to entire dense sub-tensor (hybrid) and is
            # applied to each batch
            mask = make_tensor(sparse_shape, dtype=torch.bool, device=device)
            # manually expand to match hybrid shape
            if hybrid:
                mask = mask.view(sparse_shape + tuple(1 for _ in range(n_hybrid_dim)))
                mask = mask.expand(sparse_shape + hybrid_shape)

            # mask will broadcast over the batch dims if present

            return dense * mask

        # note: order is important here, the hybrid-ness decides the inner content check which is used to build the
        # batched checker (if needed)
        check_content = _check_against_scipy_matrix
        if hybrid:
            check_content = _check_hybrid_matrix
        if batched:
            check_content = functools.partial(_check_batched, check_batch=check_content)

        sparse_sizes = [(6, 10), (0, 10), (6, 0), (0, 0)]
        blocksizes = [(2, 2), (1, 1), (1, 2)] if layout in blocked_layouts else [()]
        batch_sizes = [(3,), (1, 3), (2, 1, 3)] if batched else [()]
        hybrid_sizes = [(4, ), (2, 2)] if hybrid else [()]

        # general cases, always run
        for sparse_shape, blocksize, batch_shape, hybrid_shape in itertools.product(
                sparse_sizes, blocksizes, batch_sizes, hybrid_sizes):
            dense = _generate_subject(sparse_shape, batch_shape, hybrid_shape)
            sparse = dense.to_sparse(layout=layout, blocksize=blocksize or None, dense_dim=len(hybrid_shape))
            check_content(sparse, dense, blocksize=blocksize, batch_shape=batch_shape, hybrid_shape=hybrid_shape)
            dense_back = sparse.to_dense()
            self.assertEqual(dense, dense_back)

        # special cases for batched tensors
        if batched:
            # batched sparse tensors need only have the same number of non-zeros in each batch not nessesarily the
            # same sparsity pattern in each batch
            sparse_shape = sparse_sizes[0]
            hybrid_shape = hybrid_sizes[0]
            batch_shape = batch_sizes[0]
            shape = batch_shape + sparse_shape + hybrid_shape
            dense = make_tensor(shape, dtype=torch.float, device=device)
            blocksize = blocksizes[0]
            # number of elements/blocks in each batch (total not nnz)
            batch_mask_shape = sparse_shape
            if layout in blocked_layouts:
                # if we are blocked the mask is genereated for the block valued elemetns
                batch_mask_shape = sparse_shape[0] // blocksize[0], sparse_shape[1] // blocksize[1]

            # random bool vector w/ length equal to max possible nnz for the sparse_shape
            mask_source = make_tensor(batch_mask_shape, dtype=torch.bool, device=device).flatten()
            n_batch = functools.reduce(operator.mul, batch_shape, 1)

            # stack random permutations of the source for each batch
            mask = torch.stack([mask_source[torch.randperm(mask_source.numel())]
                               for _ in range(n_batch)], dim=0).reshape(batch_shape + batch_mask_shape)
            if layout in blocked_layouts:
                # for blocked we need to do a bit of extra work to expand the mask from blocked-space to element-space
                mask_shape = mask.shape
                mask = mask.view(mask_shape + (1, 1))
                mask = mask.expand(mask_shape + blocksize)
                mask = mask.transpose(-3, -2)
                mask = mask.flatten(-4, -3).flatten(-2, -1)
            mask_shape = mask.shape
            mask = mask.view(mask_shape + (1,) * len(hybrid_shape))
            mask = mask.expand(mask_shape + hybrid_shape)
            dense = dense * mask
            sparse = dense.to_sparse(layout=layout, blocksize=blocksize or None, dense_dim=len(hybrid_shape))
            check_content(sparse, dense, blocksize=blocksize, batch_shape=batch_shape, hybrid_shape=hybrid_shape)

            dense_back = sparse.to_dense()
            self.assertEqual(dense, dense_back)

            # if batches have different nnz we expect the conversion to throw
            mask_0 = mask[0]
            mask_1 = mask[0].clone().fill_(True)
            mask_2 = mask[0].clone().fill_(False)
            mask_true = mask_source.clone().fill_(True)
            mask_false = mask_source.clone().fill_(False)
            mask = torch.stack([(mask_0, mask_1, mask_2)[i % 3] for i in range(n_batch)], dim=0).reshape(batch_shape + mask_0.shape)
            dense = make_tensor(shape, dtype=torch.float, device=device)
            dense = dense * mask
            msg = "Expect the same number of specified elements per batch."
            with self.assertRaisesRegex(RuntimeError, msg):
                dense.to_sparse(layout=layout, blocksize=blocksize or None)

            # Should throw if there is a zero in the batch size
            dense = make_tensor((0,) + shape, dtype=torch.float, device=device)
            layout_code = str(layout).split("_")[-1]
            msg = f"to_sparse_{layout_code}: Expected product of batch dimensions to be non-zero."
            with self.assertRaisesRegex(RuntimeError, msg):
                dense.to_sparse(layout=layout, blocksize=blocksize or None)

    @skipMeta
    @all_sparse_compressed_layouts()
    @coalescedonoff
    @dtypes(torch.double)
    @unittest.skipIf(not TEST_SCIPY, "SciPy not found")
    def test_sparse_to_sparse_compressed(self, device, dtype, coalesced, layout):
        """
        This test tests conversion from COO to CSR and CSC and CSC to CSR and CSC
        by comparing to SciPy's implementation.

        Here we test only those conversion combinations that SciPy
        supports to ensure that PyTorch conversions are in the same
        page with SciPy.  Independent from SciPy, all conversion
        combinations are tested in TestSparseAny.test_to_sparse.
        """

        blocksize_kw = {}
        if layout in (torch.sparse_bsc, torch.sparse_bsr):
            blocksize_kw['blocksize'] = (2, 2)
            # block modes don't support 0 width/height
            shapes = [(6, 10)]
        elif layout in (torch.sparse_csc, torch.sparse_csr):
            shapes = [(0, 10), (6, 0), (6, 10), (0, 0)]
        else:
            raise NotImplementedError("unhandled layout")

        if layout in (torch.sparse_bsc, torch.sparse_csc):
            compressed_indices_mth = torch.Tensor.ccol_indices
            plain_indices_mth = torch.Tensor.row_indices
        elif layout in (torch.sparse_bsr, torch.sparse_csr):
            compressed_indices_mth = torch.Tensor.crow_indices
            plain_indices_mth = torch.Tensor.col_indices
        else:
            raise NotImplementedError("unhandled layout")

        for shape in shapes:
            sparse_dim = 2
            nnz = shape[0] * shape[1] // 2
            sparse, _, _ = self.genSparseTensor(shape, sparse_dim, nnz, coalesced, device, dtype)
            sp_matrix = self._construct_sp_matrix(sparse, layout)
            pt_matrix = sparse.to_sparse(layout=layout, **blocksize_kw)

            self.assertEqual(layout, pt_matrix.layout)
            self.assertEqual(sp_matrix.shape, pt_matrix.shape)
            self.assertEqual(torch.tensor(sp_matrix.indptr, dtype=torch.int64), compressed_indices_mth(pt_matrix))
            self.assertEqual(torch.tensor(sp_matrix.indices, dtype=torch.int64), plain_indices_mth(pt_matrix))
            self.assertEqual(torch.tensor(sp_matrix.data), pt_matrix.values())

            sparse_csc = sparse.to_sparse_csc()
            sp_matrix = self._construct_sp_matrix(sparse_csc, layout)
            pt_matrix = sparse_csc.to_sparse(layout=layout, **blocksize_kw)

            self.assertEqual(layout, pt_matrix.layout)
            self.assertEqual(sp_matrix.shape, pt_matrix.shape)
            self.assertEqual(torch.tensor(sp_matrix.indptr, dtype=torch.int64), compressed_indices_mth(pt_matrix))
            self.assertEqual(torch.tensor(sp_matrix.indices, dtype=torch.int64), plain_indices_mth(pt_matrix))
            self.assertEqual(torch.tensor(sp_matrix.data), pt_matrix.values())

    @unittest.skipIf(not TEST_CUDA_CUDSS, "The test requires cudss")
    @dtypes(*floating_types())
    def test_linalg_solve_sparse_csr_cusolver(self, device, dtype):
        # https://github.com/krshrimali/pytorch/blob/f5ee21dd87a7c5e67ba03bfd77ea22246cabdf0b/test/test_sparse_csr.py

        try:
            spd = torch.rand(4, 3)
            A = spd.T @ spd
            b = torch.rand(3).cuda()
            A = A.to_sparse_csr().cuda()
            x = torch.sparse.spsolve(A, b)
        except RuntimeError as e:
            if "Calling linear solver with sparse tensors requires compiling " in str(e):
                self.skipTest("PyTorch was not built with cuDSS support")

        samples = sample_inputs_linalg_solve(None, device, dtype)

        for sample in samples:
            if sample.input.ndim != 2:
                continue

            out = torch.zeros(sample.args[0].size(), dtype=dtype, device=device)
            if sample.args[0].ndim != 1 and sample.args[0].size(-1) != 1:
                with self.assertRaisesRegex(RuntimeError, "b must be a 1D tensor"):
                    out = torch.linalg.solve(sample.input.to_sparse_csr(), *sample.args, **sample.kwargs)
                break
            if not sample.args[0].numel():
                with self.assertRaisesRegex(RuntimeError,
                                            "Expected non-empty other tensor, but found empty tensor"):
                    torch.linalg.solve(sample.input.to_sparse_csr(), *sample.args, **sample.kwargs, out=out)
                break

            expect = torch.linalg.solve(sample.input, *sample.args, **sample.kwargs)
            sample.input = sample.input.to_sparse_csr()
            if sample.args[0].ndim != 1 and sample.args[0].size(-1) == 1:
                expect = expect.squeeze(-1)
                sample.args = (sample.args[0].squeeze(-1), )
            out = torch.linalg.solve(sample.input, *sample.args, **sample.kwargs)
            self.assertEqual(expect, out)


def skipIfNoTriton(cls):
    from torch.utils._triton import has_triton

    # no-op if triton is present
    if has_triton():
        return cls
    else:

        @functools.wraps(cls, updated=())
        class skipped_cls(cls):
            def setUp(self):
                self.skipTest("Triton is not available.")

        return skipped_cls

@skipIfNoTriton
class TestSparseCompressedTritonKernels(TestCase):

    def _to_block_triangular_inplace(self, d, row_block, col_block):
        """
        This function modifies `d` to become (upper/lower) block-triangular in-place.
        It is assumed that `d.shape[-2]` is divisible by `row_block` and
        `d.shape[-1]` is divisible by `col_block`.
        """

        from torch.sparse._triton_ops import tile_to_blocksize

        m, n = d.shape[-2:]
        d_tiled = tile_to_blocksize(d, (row_block, col_block))
        d_tiled = d_tiled.moveaxis(-4, -1).moveaxis(-4, -1)
        if m // row_block > n // col_block:
            d_tiled.tril_()
        else:
            d_tiled.triu_()

        return d

    @onlyCUDA
    @skipIfRocm(msg="test is too slow on ROCm stack")
    @dtypes(torch.half, torch.bfloat16, torch.float)
    @dtypesIfCUDA(torch.half, *[torch.bfloat16] if SM80OrLater else [], torch.float)
    @unittest.skipIf(IS_FBCODE and IS_REMOTE_GPU, "Test requires Triton")
    def test_triton_bsr_softmax(self, device, dtype):
        from functools import partial
        from torch.sparse._triton_ops import bsr_softmax

        tensor = partial(make_tensor, device=device, dtype=dtype, low=1.0, high=3.0)

        # NOTE: batch dims with zero sizes are not supported in `to_sparse_bsr`.
        batches = [(), (2,), (2, 2)]
        size = [6, 12, 0]
        block_size = [2, 3]

        # General correctness
        for row_block, col_block, b, m, n in itertools.product(block_size, block_size, batches, size, size):
            input = tensor(b + (m, n))
            input.diagonal(dim1=-2, dim2=-1).fill_(m * n)
            input = self._to_block_triangular_inplace(input, row_block, col_block)

            bsr = input.to_sparse_bsr((row_block, col_block))
            coo = input.to_sparse().to(torch.float)

            res_tri = bsr_softmax(bsr)
            res_coo = torch.sparse.softmax(coo, -1)
            self.assertEqual(res_tri, res_coo.to(input.dtype))

        # Test long rows which exceed Triton's max numel limit set to 2 ** 17
        input = tensor(b + (1, 150000))
        bsr = input.to_sparse_bsr(1)
        self.assertEqual(input.softmax(-1), bsr_softmax(bsr))

    @parametrize("block_size", [16, 32, 64])
    @parametrize("index_dtype", [torch.int32, torch.int64])
    @onlyCUDA
    @dtypes(torch.half, torch.bfloat16, torch.float)
    @dtypesIfCUDA(torch.half, *[torch.bfloat16] if SM80OrLater else [], torch.float)
    @unittest.skipIf((not TEST_WITH_TORCHINDUCTOR) or (IS_FBCODE and IS_REMOTE_GPU) or torch._running_with_deploy(),
                     "Skipped for deploy and internal with remote GPUs")
    def test_triton_bsr_dense_bmm(self, device, dtype, index_dtype, block_size):
        from functools import partial
        from torch.sparse._triton_ops import bsr_dense_mm

        def kernel_impl(*args, **kwargs):
            return bsr_dense_mm(*args, skip_checks=True, **kwargs)

        kernel = torch._TritonLibrary.registerOp(
            "_triton_bsr_dense_mm_out",
            "_triton_bsr_dense_mm_out(Tensor bsr, Tensor dense, *, Tensor(a!) out) -> Tensor(a!)",
            kernel_impl,
            "SparseCsrCUDA"
        )

        # kernel != kernel_impl means dispatch was already registered.
        # This is exactly what we need!
        self.assertTrue(kernel is not kernel_impl)

        # Note that each value in a non-zero block is in range block_size * [low^2, high^2).
        tensor = partial(make_tensor, device=device, dtype=dtype, low=0.5, high=1.5)

        # NOTE: batch dims with zero sizes are not supported in `to_sparse_bsr`.
        batches = [(), (2,), (2, 2)]
        size = [128, 256, 0]

        # Whether to make inputs orthogonal so that the product is zero
        make_orthogonal = [True, False]

        for bd, bs, m, n, k, is_ortho in itertools.product(batches, batches, size, size, size, make_orthogonal):
            bsr = tensor(bs + (m, k))
            # NOTE: do not get confused, it will be transposed
            dense = tensor(bd + (n, k))

            if is_ortho:
                bsr = torch.cat((bsr, torch.zeros_like(bsr)), dim=-1)
                dense = torch.cat((torch.zeros_like(dense), dense), dim=-1)

            bsr = bsr.to_sparse_bsr(block_size)

            if bsr.dim() == 2 and dtype != torch.float:
                # Test against linear to check dispatch
                # which takes place for torch.half and torch.bfloat16.
                res_dense = torch.nn.functional.linear(dense, bsr.to_dense())
                res_tri_out = torch.empty_like(res_dense)
                res_tri = torch.nn.functional.linear(dense, bsr, out=res_tri_out)

                # Check dispatch worked with non-trivial outputs
                if m > 0 and n > 0 and k > 0:
                    self.assertTrue(kernel.kernel_invoked)
                    kernel.kernel_invoked = False
            else:
                # Otherwise check correctness against bmm
                # since nn.linear does not support bsr.dim() > 2.
                res_dense = bsr.to_dense() @ dense.transpose(-2, -1)
                res_tri_out = torch.empty_like(res_dense)
                res_tri = kernel(bsr, dense.transpose(-2, -1), out=res_tri_out)

            self.assertTrue(res_tri is res_tri_out)
            self.assertEqual(res_tri, res_dense)

            res_dense = bsr.to_dense() @ dense.transpose(-2, -1)
            # check whether bsr_dense_mm handles different grid sizes
            # None means max possible grid size which is CUDA-dependent.
            grid_size = (None, 2, 4)
            grid_gen = itertools.product(grid_size, repeat=3)
            for grid in grid_gen:
                res_tri = torch.sparse._triton_ops.bsr_dense_mm(
                    bsr,
                    dense.transpose(-2, -1),
                    max_grid=grid,
                )
                self.assertEqual(res_tri, res_dense)

    @onlyCUDA
    @dtypes(torch.half)
    @unittest.skipIf(IS_FBCODE and IS_REMOTE_GPU or torch._running_with_deploy(),
                     "Skipped for deploy and internal with remote GPUs")
    def test_triton_bsr_dense_bmm_error_messages(self, device, dtype):
        from torch.sparse._triton_ops import bsr_dense_mm

        rhs = torch.rand(32, 32, dtype=dtype, device=device)
        lhs = rhs.to_sparse_bsr(16)
        with self.assertRaisesRegex(ValueError, "only BSR sparse format is supported"):
            bsr_dense_mm(lhs.to_sparse_bsc(16), rhs)
        with self.assertRaisesRegex(ValueError, "on the same GPU device"):
            bsr_dense_mm(lhs, rhs.cpu())
        if torch.cuda.device_count() > 1:
            with self.assertRaisesRegex(ValueError, "on the same GPU device"):
                bsr_dense_mm(lhs.to("cuda:0"), rhs.to("cuda:1"))
        with self.assertRaisesRegex(ValueError, "all inputs are expected to be of the same dtype"):
            bsr_dense_mm(lhs, rhs.to(torch.float))
        with self.assertRaisesRegex(ValueError, r"and one of \(half, bfloat16, float32\)"):
            bsr_dense_mm(lhs.to(torch.double), rhs.to(torch.double))
        with self.assertRaisesRegex(ValueError, "all inputs involved in the matrix product are expected to be at least 2D"):
            bsr_dense_mm(lhs, torch.rand(1, dtype=dtype, device=device))
        with self.assertRaisesRegex(ValueError,
                                    "sizes involved in the matrix product are not compatible for matrix multiplication"):
            bsr_dense_mm(lhs, torch.rand(1, 1, dtype=dtype, device=device))
        with self.assertRaisesRegex(ValueError,
                                    r"dense.size\(-1\) == 15 should be divisible by 16"):
            bsr_dense_mm(lhs, torch.rand(32, 15, dtype=dtype, device=device))
        # Blocksizes check
        for blocksize in (15, 30):
            n = blocksize * 2
            rhs = torch.rand(n, n, dtype=dtype, device=device)
            lhs = rhs.to_sparse_bsr(blocksize)
            with self.assertRaisesRegex(ValueError, "should be at least 16 and a power of 2"):
                bsr_dense_mm(lhs, rhs)
        # out check
        rhs = torch.rand(2, 32, 32, dtype=dtype, device=device)
        lhs = rhs.to_sparse_bsr(16)
        with self.assertRaisesRegex(ValueError, r"`out` argument has wrong shape"):
            out = torch.rand(2, 30, 30, dtype=dtype, device=device)
            bsr_dense_mm(lhs, rhs, out=out)
        with self.assertRaisesRegex(ValueError, r"only row-major/col-major `out`"):
            out = torch.rand(32, 32, 2, dtype=dtype, device=device).transpose(0, -1)
            bsr_dense_mm(lhs, rhs, out=out)

    @parametrize("block_size", [16, 32, 64])
    @onlyCUDA
    @skipIfRocm
    @dtypes(torch.half, torch.bfloat16, torch.float)
    @dtypesIfCUDA(torch.half, *[torch.bfloat16] if SM80OrLater else [], torch.float)
    @unittest.skipIf(IS_FBCODE and IS_REMOTE_GPU, "Test requires Triton")
    @precisionOverride({torch.float16: 1e-3})
    def test_triton_scaled_dot_product_attention(self, device, dtype, block_size):
        from functools import partial
        from torch.sparse._triton_ops import _scaled_dot_product_attention

        # Note that each value in a non-zero block is in range block_size * [low^2, high^2).
        tensor = partial(make_tensor, device=device, dtype=dtype, low=0.3, high=1.2)

        def broadcast_input(*ts):
            batch_dims = torch.broadcast_shapes(*(t.shape[:-2] for t in ts))
            yield from (torch.broadcast_to(t, batch_dims + t.shape[-2:]) for t in ts)

        # NOTE: batch dims with zero sizes are not supported in `to_sparse_bsr`.
        batches = [(), (2,), (2, 2)]
        size = [128, 256, 0]

        for bam, bq, bk, bv, m, n, k in itertools.product(batches, batches, batches, batches, size, size, size):
            query = tensor(bq + (m, k))
            key = tensor(bk + (n, k))
            value = tensor(bv + (n, k))

            # We make attn_mask block lower/upper triangular so that BSR and Strided
            # function variants are directly comparable.
            attn_mask = torch.ones(bam + (m, n), device=device, dtype=torch.bool)
            attn_mask = self._to_block_triangular_inplace(attn_mask, block_size, block_size)
            attn_mask_bsr = attn_mask.to_sparse_bsr(block_size)

            # NOTE: only boolean mask is directly compatible with the Strided version
            # without any pre-/post-processing. Hence we test against a boolean mask.
            for scale in (None, 1. / 16):
                if scale is None and query.size(-1) == 0:
                    scale = 1
                expected = torch.nn.functional.scaled_dot_product_attention(
                    *broadcast_input(query, key, value, attn_mask), scale=scale
                )

                for mask_dtype in (torch.bool, dtype):
                    res = _scaled_dot_product_attention(query, key, value, attn_mask_bsr.to(mask_dtype), scale=scale)
                    self.assertEqual(res, expected)


    @parametrize("block_size", [16, 32, 64])
    @onlyCUDA
    @skipIfRocm
    @dtypes(torch.half, torch.bfloat16, torch.float)
    @dtypesIfCUDA(torch.half, *[torch.bfloat16] if SM80OrLater else [], torch.float)
    @unittest.skipIf(IS_FBCODE and IS_REMOTE_GPU, "Test requires Triton")
    def test_triton_sampled_addmm(self, device, dtype, block_size):
        from functools import partial
        from torch.sparse._triton_ops import sampled_addmm, broadcast_batch_dims_bsr

        # Note that each value in a non-zero block is in range block_size * [low^2, high^2).
        tensor = partial(make_tensor, device=device, dtype=dtype, low=0.3, high=1.2)

        # NOTE: batch dims with zero sizes are not supported in `to_sparse_bsr`.
        batches = [(), (2,), (2, 2)]
        size = [128, 256, 0]

        delta_k = (-3,)
        for bi, bm1, bm2, m, n, k, dk in itertools.product(batches, batches, batches, size, size, size, delta_k):
            # Test not powers of 2 ks as well.
            k = max(0, k + dk)
            # Non-trivial sparsity pattern.
            # Plus with tril inputs the result is also tril,
            # so we can compare BSR and CSR implementations.
            input = tensor(bi + (m, n)).tril_()
            bsr = input.to_sparse_bsr(block_size)
            mat1 = tensor(bm1 + (m, k)).tril_()
            mat2 = tensor(bm2 + (k, n)).tril_()

            batch_dim = torch.broadcast_shapes(input.shape[:-2], mat1.shape[:-2], mat2.shape[:-2])

            csr = input.broadcast_to(batch_dim + input.shape[-2:]).to_sparse_csr().to(torch.float)
            mat1csr = mat1.broadcast_to(batch_dim + mat1.shape[-2:]).to(torch.float)
            mat2csr = mat2.broadcast_to(batch_dim + mat2.shape[-2:]).to(torch.float)

            input_broadcasted_clone = broadcast_batch_dims_bsr(
                "test_triton_sampled_addmm",
                bsr, mat1, mat2
            ).clone()
            input_broadcasted_clone = torch.sparse_compressed_tensor(
                input_broadcasted_clone.crow_indices(),
                input_broadcasted_clone.col_indices(),
                # For testing `out=` let's make values to have "weird" strides
                # so that if the kernel modifies values to it's needs, the result
                # is being compied into out.values.
                input_broadcasted_clone.values().transpose(-3, -2).contiguous().transpose(-3, -2),
                layout=input_broadcasted_clone.layout,
                size=input_broadcasted_clone.shape
            )

            scalars = (0.0, 2.0)
            for alpha, beta, out in itertools.product(scalars, scalars, (None, input_broadcasted_clone)):
                res_tri = sampled_addmm(bsr, mat1, mat2, alpha=alpha, beta=beta, out=out)
                if out is not None:
                    self.assertTrue(res_tri is out)

                batch_broadcasted_shape = torch.broadcast_shapes(*(t.shape[:-2] for t in (input, mat1, mat2)))
                self.assertTrue(res_tri.shape == batch_broadcasted_shape + (m, n))

                res_csr = torch.sparse.sampled_addmm(csr, mat1csr, mat2csr, alpha=alpha, beta=beta).to(input.dtype)
                self.assertEqual(res_tri.to_dense(), res_csr.to_dense())

                # Check different grid sizes to make sure that input slicing works
                # if this input is larger than the grid.
                grid_size = (3, None)
                grid_gen = itertools.product(grid_size, repeat=2)
                for grid in grid_gen:
                    res_tri_grid = sampled_addmm(bsr, mat1, mat2, alpha=alpha, beta=beta, max_grid=grid)
                    self.assertEqual(res_tri, res_tri_grid)

    @onlyCUDA
    @skipIfRocm
    @dtypes(torch.half, torch.bfloat16, torch.float)
    @dtypesIfCUDA(torch.half, *[torch.bfloat16] if SM80OrLater else [], torch.float)
    @unittest.skipIf(IS_FBCODE and IS_REMOTE_GPU, "Test requires Triton")
    def test_triton_scatter_mm(self, device, dtype):
        from torch.sparse._triton_ops import scatter_mm
        from functools import partial
        tensor = partial(make_tensor, device=device, dtype=dtype, low=0.5, high=1.5)
        sizes = [8, 16]
        for m, k, n in itertools.product(sizes, sizes, sizes):
            blocks = torch.stack([tensor(m, k), tensor(m, k)])
            others = torch.stack([tensor(k, n), tensor(k, n)])

            expected = torch.stack([blocks[0] @ others[0] + blocks[1] @ others[0],
                                    blocks[0] @ others[1],
                                    blocks[1] @ others[1]])

            indices_data = (
                'scatter_mm',
                torch.tensor([0, 2, 3, 4], dtype=torch.int32, device=device),
                torch.tensor([[0, 0], [1, 0], [0, 1], [1, 1]], dtype=torch.int32, device=device))

            result = scatter_mm(blocks, others, indices_data=indices_data)

            self.assertEqual(result, expected)

            indices_data = (
                'bsr_strided_mm',
                torch.tensor([0, 2, 4, 5, 6], dtype=torch.int32, device=device),
                torch.tensor([0, n, 2 * n * m, 2 * n * m + n], dtype=torch.int32, device=device),
                torch.tensor([1, 0, 1, 0, 1, 1], dtype=torch.int32, device=device),
                torch.tensor([0, 2 * k * n, n, 2 * k * n + n, 2 * k * n, 2 * k * n + n],
                             dtype=torch.int32, device=device),
                dict(SPLIT_N=2, is_compressed=False, TILE_M=m, TILE_N=n, GROUP_SIZE=1)
            )

            for bsize in [(), (2,), (3, 4)]:
                other = tensor(*bsize, 2 * k, 2 * n)
                expected = torch.cat([
                    torch.cat([blocks[1], blocks[0]], dim=1),
                    torch.cat([torch.zeros_like(blocks[0]), blocks[1]], dim=1)], dim=0) @ other
                result = scatter_mm(blocks, other, indices_data=indices_data)
                self.assertEqual(result, expected)

    @parametrize("blocksize", [2, '2x3', 16, '16x32', 32, 64])
    @onlyCUDA
    @dtypes(torch.half, torch.bfloat16, torch.float)
    @dtypesIfCUDA(torch.half, *[torch.bfloat16] if SM80OrLater else [], torch.float)
    @unittest.skipIf(IS_FBCODE and IS_REMOTE_GPU, "Test requires Triton")
    def test_triton_bsr_scatter_mm(self, device, dtype, blocksize):
        import triton
        from torch.sparse._triton_ops import bsr_scatter_mm, bsr_scatter_mm_indices_data
        from functools import partial
        if isinstance(blocksize, str):
            blocksize = tuple(map(int, blocksize.split('x')))
        else:
            blocksize = (blocksize,) * 2
        # Note that each value in a non-zero block is in range blocksize * [low^2, high^2).
        tensor = partial(make_tensor, device=device, dtype=dtype, low=0.5, high=1.5)

        # NOTE: batch dims with zero sizes are not supported in `to_sparse_bsr`.
        batches = [(), (2,), (2, 2)]
        sizes = [blocksize[0], 2 * blocksize[0], 4 * blocksize[0]]
        sizes_K = [blocksize[1], 2 * blocksize[1]]

        for bd, bs, M, K, N, has_zero_row_block in itertools.product(batches, batches[:1], sizes, sizes_K, sizes, (False, True)):
            bsr_dense = tensor(bs + (M, K))
            if has_zero_row_block:
                if M > blocksize[0]:
                    bsr_dense[:blocksize[0]].zero_()
                else:
                    continue
            bsr = bsr_dense.to_sparse_bsr(blocksize)
            dense = tensor(bd + (K, N))
            expected = bsr.to_dense() @ dense

            for indices_format in ('bsr_strided_mm', 'bsr_strided_mm_compressed', 'scatter_mm'):
                if indices_format in {'bsr_strided_mm', 'bsr_strided_mm_compressed'}:
                    SPLIT_N_list = [N]
                    while SPLIT_N_list[-1] > 1:
                        SPLIT_N_list.append(max(1, SPLIT_N_list[-1] // 2))
                else:
                    SPLIT_N_list = [1]
                for SPLIT_N in SPLIT_N_list:
                    indices_data = bsr_scatter_mm_indices_data(
                        bsr, dense, indices_format=indices_format, SPLIT_N=SPLIT_N)
                    try:
                        result = bsr_scatter_mm(bsr, dense, indices_data=indices_data)
                    except triton.compiler.OutOfResources:
                        # ensure that there was at least one succesful test:
                        assert SPLIT_N < SPLIT_N_list[0]
                        break

                    self.assertEqual(result, expected)
        torch.sparse._triton_ops._bsr_scatter_mm_indices_data.cache_clear()

    def test_TensorAsKey(self, device):
        from torch.sparse._triton_ops import TensorAsKey
        assertEqualOptions = dict(exact_dtype=True, exact_device=True, exact_layout=True)

        t = torch.tensor([1, 2, 3, 4], dtype=torch.int64, device=device)
        key = TensorAsKey(t)
        self.assertTrue(key == TensorAsKey(t))
        self.assertTrue(key.obj is t)

        t2 = t[:]
        key2 = TensorAsKey(t2)
        self.assertTrue(key == key2)
        self.assertEqual(key2.obj, t, **assertEqualOptions)
        # deleting object leads to dead key
        del t2
        self.assertTrue(key2.obj is None)
        self.assertTrue(key.obj is t)

        # key with different storage offset and shape:
        self.assertFalse(key == TensorAsKey(t[1:]))

        # key with different strides:
        self.assertFalse(key == TensorAsKey(t[::2]))

        # when object dies, make sure that key represents a dead
        # object as well:
        del t
        self.assertTrue(key.obj is None)

        # Storing a tensor as a dict key:
        d = {}
        t3 = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=device)
        key3 = TensorAsKey(t3)
        d[key3] = 123
        self.assertTrue(d.get(key3) == 123)
        t3_ = t3[:]
        self.assertTrue(d.get(TensorAsKey(t3_)) == 123)
        self.assertTrue(d.get(TensorAsKey(t3.clone())) is None)

        d[TensorAsKey(t3_)] = 567
        self.assertTrue(d.get(key3) == 567)

        # t3 and t3_ reference the same data, so, the key becomes dead
        # (that is, its .obj property returns None) until all
        # references are deleted:
        del t3
        self.assertTrue(key3.obj is not None)
        self.assertTrue(d.get(key3) == 567)
        del t3_
        self.assertTrue(key3.obj is None)
        self.assertTrue(d.get(key3) == 567)

        # Storing a tensor as a dict key and value:
        d = {}
        t4 = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=device)
        key4 = TensorAsKey(t4)
        d[key4] = (t4, 123)
        self.assertEqual(d.get(key4), (t4, 123), **assertEqualOptions)
        # when object is deleted, the key represents an alive object
        # because the object is referenced by the dict item value:
        del t4
        self.assertTrue(key4.obj is not None)
        # This also means that the life time of the tensor is same as
        # the life time of the corresponding dict item:
        del d[key4]
        self.assertTrue(key4.obj is None)

        # Storing a tensor as a dict key and value wrapped with TensorAsKey:
        d = {}
        t5 = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=device)
        key5 = TensorAsKey(t5)
        d[key5] = (key5, 567)
        self.assertEqual(d.get(key5), (key5, 567), **assertEqualOptions)
        self.assertTrue(key5.obj is not None)
        # when object is deleted, it will be dead as the wrapped value
        # hold the tensor instance as a weakref:
        del t5
        self.assertTrue(key5.obj is None)
        # but key is still valid:
        self.assertEqual(d.get(key5), (key5, 567), **assertEqualOptions)

    @suppress_warnings
    @parametrize("op", ['bsr_dense_addmm', 'bsr_dense_mm', 'bsr_dense_linear', '_int_bsr_dense_addmm'])
    @parametrize("blocksize", [16, '16x32', 32])
    @onlyCUDA
    @skipIfRocm
    @dtypes(torch.half, torch.bfloat16, torch.float, torch.int8)
    @dtypesIfCUDA(torch.half, *[torch.bfloat16] if SM80OrLater else [], torch.float, torch.int8)
    @unittest.skipIf(IS_FBCODE and IS_REMOTE_GPU, "Test requires Triton")
    def test_triton_kernel(self, op, device, dtype, blocksize):
        from torch.sparse._triton_ops import bsr_dense_addmm, bsr_dense_mm, _int_bsr_dense_addmm
        from torch.sparse._triton_ops_meta import (create_blocked_tensor, get_meta,
                                                   optimize_bsr_dense_addmm, dump)

        def bsr_dense_linear(input, weights, bias=None):
            return torch.nn.functional.linear(input, weights, bias=bias).transpose(-1, -2)

        operation = dict(bsr_dense_addmm=bsr_dense_addmm, bsr_dense_mm=bsr_dense_mm, bsr_dense_linear=bsr_dense_linear,
                         _int_bsr_dense_addmm=_int_bsr_dense_addmm)[op]

        def reference(input, mat1, mat2, beta=1, alpha=1, op=op):
            assert mat1.layout is torch.strided
            assert mat2.layout is torch.strided
            if dtype is torch.int8:
                if op == '_int_bsr_dense_addmm':
                    return beta * input + alpha * torch._int_mm(mat1, mat2)
                # workaround RuntimeError: "addmm_cuda" not implemented for 'Char'
                return beta * input + alpha * torch._int_mm(mat1, mat2).to(torch.int8)
            return beta * input + alpha * (mat1 @ mat2)

        if op == '_int_bsr_dense_addmm':
            # _int_bsr_dense_addmm is same as bsr_dense_addmm except
            # with int8 inputs, _int_bsr_dense_addmm returns int32
            # result. This is covered by operation and reference
            # definitions above and all other definitions below are
            # identical between _int_bsr_dense_addmm and
            # bsr_dense_addmm.
            op = 'bsr_dense_addmm'

        def nc_copy(t, axes=(-1,)):
            """Return a copy of input.

            The returned copy will be a non-contiguous tensor.
            """
            if t.layout is torch.strided:
                shape = list(t.shape)
                for a in axes:
                    shape[a] *= 2
                r = torch.empty(shape, dtype=t.dtype, device=t.device)
                s = r[tuple(slice(None, None, 2 if t.shape[i] != r.shape[i] else None) for i in range(t.ndim))]
                s.copy_(t)
                return s
            elif t.layout is torch.sparse_bsr:
                compressed_indices = t.crow_indices()
                plain_indices = t.col_indices()
                return torch.sparse_compressed_tensor(compressed_indices, plain_indices, nc_copy(t.values()),
                                                      t.shape, layout=t.layout)
            else:
                raise NotImplementedError(t.layout)

        if isinstance(blocksize, str):
            BM, BK = tuple(map(int, blocksize.split('x')))
        else:
            BM, BK = (blocksize,) * 2

        if op in {"bsr_dense_linear"} and BM != BK:
            # todo: eliminate this skip
            self.skipTest(f"{op} does not support non-square blocks")

        if op in {"bsr_dense_linear"} and dtype is torch.int8:
            # todo: eliminate this skip
            self.skipTest(f"{op} does not support int8")

        if dtype is torch.int8 and min(BM, BK) < 32:
            self.skipTest("triton kernel does not support support int8 blocks smaller than 32")

        beta_lst = dict(bsr_dense_addmm=[0, 1, 2], bsr_dense_mm=[0], bsr_dense_linear=[1])[op]
        alpha_lst = dict(bsr_dense_addmm=[0, 1, 2], bsr_dense_mm=[1], bsr_dense_linear=[1])[op]
        sparsity_lst = [0, 0.5, 1]
        blocks_per_row_lst = [1, 2]
        blocks_per_col_lst = [1, 2]
        result_cols_lst = [16, 32, 64]
        for beta, alpha, sparsity, blocks_per_row, blocks_per_col, N in itertools.product(
                beta_lst, alpha_lst, sparsity_lst, blocks_per_row_lst, blocks_per_col_lst, result_cols_lst):
            M = BM * blocks_per_row
            K = BK * blocks_per_col
            mat1 = create_blocked_tensor(0, M, K, (BM, BK), sparsity, dtype, device=device)
            bsr = mat1.to_sparse_bsr((BM, BK))
            mat2 = make_tensor(K, N, dtype=dtype, device=device, low=0.5, high=1.5)
            input = make_tensor(M, N, dtype=dtype, device=device, low=0.5, high=1.5)

            if 0 and op == "bsr_dense_addmm":
                # Find optimal kernel parameters, the speed-up is
                # about 10x for running this test.
                #
                # Enable this if-block when the test method is
                # updated, run the test, and finally, disable the
                # if-block.
                key = (M, K, N, BM, BK, beta == 0, beta == 1, alpha == 1)
                meta = get_meta(op, key, version=(0, dtype, 0.5))
                if meta is None:
                    optimize_bsr_dense_addmm(M, K, N, BM, BK, beta=beta, alpha=alpha, dtype=dtype, sparsity=0.5)
                    meta = get_meta(op, key, version=(0, dtype, 0.5))
                    assert meta is not None
                    dump()  # this will update torch/sparse/_triton_ops_meta.py

            expected = reference(input, mat1, mat2, beta=beta, alpha=alpha)
            kwargs = dict(bsr_dense_addmm=dict(beta=beta, alpha=alpha), bsr_dense_mm={},
                          bsr_dense_linear=dict(bias=input.transpose(-1, -2)))[op]

            args = dict(bsr_dense_addmm=(input, bsr, mat2), bsr_dense_mm=(bsr, mat2),
                        bsr_dense_linear=(mat2.transpose(-1, -2), bsr))[op]
            result = operation(*args, **kwargs)
            self.assertEqual(result, expected)

            # Test non-contiguous input tensors:
            nc_mat2 = nc_copy(mat2)
            nc_input = nc_copy(input)
            nc_bsr = nc_copy(bsr)

            args = dict(bsr_dense_addmm=(input, bsr, nc_mat2), bsr_dense_mm=(bsr, nc_mat2),
                        bsr_dense_linear=(nc_mat2.transpose(-1, -2), bsr))[op]
            result = operation(*args, **kwargs)
            self.assertEqual(result, expected)

            # todo: add bsr_dense_linear to the set below (currently,
            # nn.linear has unnecessarily restrictive arguments
            # checks).
            if op in {'bsr_dense_addmm', 'bsr_dense_mm'}:
                args = dict(bsr_dense_addmm=(input, nc_bsr, mat2), bsr_dense_mm=(nc_bsr, mat2),
                            bsr_dense_linear=(mat2.transpose(-1, -2), nc_bsr))[op]
                result = operation(*args, **kwargs)
                self.assertEqual(result, expected)

            if op in {'bsr_dense_addmm', 'bsr_dense_linear'}:
                args = dict(bsr_dense_addmm=(nc_input, bsr, nc_mat2),
                            bsr_dense_linear=(nc_mat2.transpose(-1, -2), bsr))[op]
                kwargs = dict(bsr_dense_addmm=dict(beta=beta, alpha=alpha),
                              bsr_dense_linear=dict(bias=nc_input.transpose(-1, -2)))[op]
                result = operation(*args, **kwargs)
                self.assertEqual(result, expected)

    @parametrize("op", ['bsr_dense_addmm', '_int_bsr_dense_addmm'])
    @onlyCUDA
    @skipIfRocm
    @dtypes(torch.half, torch.bfloat16, torch.float, torch.int8)
    @dtypesIfCUDA(torch.half, *[torch.bfloat16] if SM80OrLater else [], torch.float, torch.int8)
    @unittest.skipIf(IS_FBCODE and IS_REMOTE_GPU, "Test requires Triton")
    def test_triton_tune(self, op, device, dtype):
        from torch.sparse._triton_ops import bsr_dense_addmm, _int_bsr_dense_addmm
        from torch.sparse._triton_ops_meta import (create_blocked_tensor, tune_bsr_dense_addmm, tune__int_bsr_dense_addmm, get_meta)

        operation = dict(bsr_dense_addmm=bsr_dense_addmm, _int_bsr_dense_addmm=_int_bsr_dense_addmm)[op]
        tuner = dict(bsr_dense_addmm=tune_bsr_dense_addmm,
                     _int_bsr_dense_addmm=tune__int_bsr_dense_addmm)[op]

        if op == '_int_bsr_dense_addmm':
            M, K, N = 32, 32, 32
            blocksize = (32, 32)
        else:
            M, K, N = 16, 16, 32
            blocksize = (16, 16)
        sparsity = 1.0
        bsr = create_blocked_tensor(0, M, K, blocksize, sparsity, dtype, device).to_sparse_bsr(blocksize)
        sparsity = 1 - bsr._nnz() * blocksize[0] * blocksize[1] / (M * K)
        input = make_tensor(K, N, dtype=dtype, device=device)
        dense = make_tensor(K, N, dtype=dtype, device=device)

        if op in {'bsr_dense_addmm', '_int_bsr_dense_addmm'}:
            args = (input, bsr, dense)

            def get_current_meta():
                version = (0, dtype, sparsity)
                meta_key = (M, K, N, *blocksize, False, True, True)
                return get_meta(op, meta_key, version=version, exact=True)
        else:
            raise NotImplementedError(op)

        self.assertEqual(get_current_meta(), None)

        meta = tuner(*args, **dict(store=True, verbose=False))
        self.assertEqual(get_current_meta(), meta)

        expected = operation(*args)
        result = operation(*args, **dict(meta=meta))
        self.assertEqual(result, expected)

    @onlyCUDA
    @skipIfRocm
    @unittest.skipIf(IS_FBCODE and IS_REMOTE_GPU, "Test requires Triton")
    def test_triton_bsr_dense_addmm_meta(self, device):
        from torch.sparse._triton_ops import bsr_dense_addmm_meta
        from torch.sparse._triton_ops_meta import update as update_bsr_dense_addmm_meta

        dtype = torch.float32
        Ms = Ks = 16
        beta = 0.0
        alpha = 1.0

        def get_meta(M, K, N, sparsity=None):
            return bsr_dense_addmm_meta(M, K, N, Ms, Ks, beta, alpha, dtype=dtype, sparsity=sparsity,
                                        _version="test_triton_bsr_dense_addmm_meta")

        def update_meta(M, K, N, value, sparsity=0.5):
            key = (M, K, N, Ms, Ks, beta == 0, beta == 1, alpha == 1)
            update_bsr_dense_addmm_meta("bsr_dense_addmm", torch.cuda.get_device_name(),
                                        ("test_triton_bsr_dense_addmm_meta", dtype, sparsity),
                                        key, value)

        def get_meta_with_checks(M, K, N, warn_count=0, sparsity=None):
            f = io.StringIO()
            with redirect_stderr(f):
                result = get_meta(M, K, N, sparsity=sparsity)
            msg = f.getvalue()
            FileCheck().check_count(
                str=f"UserWarning: bsr_dense_addmm uses non-optimal triton kernel parameters for M={M} K={K} N={N}",
                count=warn_count, exactly=True
            ).run(msg)
            return result

        # Test warn_once when requesting non-existing tuned parameters multiple times
        f = io.StringIO()
        with redirect_stderr(f):
            for i in range(5):
                get_meta(16, 16, 16)
            for i in range(5):
                get_meta(16, 16, 32)

        msg = f.getvalue()
        FileCheck().check_count(
            str="UserWarning: bsr_dense_addmm uses non-optimal triton kernel parameters for M=16 K=16 N=16", count=1, exactly=True
        ).run(msg)
        FileCheck().check_count(
            str="UserWarning: bsr_dense_addmm uses non-optimal triton kernel parameters for M=16 K=16 N=32", count=1, exactly=True
        ).run(msg)

        # Test warn_once when tuned parameters are missing
        default_meta = dict(GROUP_SIZE_ROW=4, SPLIT_N=2, num_stages=1, num_warps=4)
        self.assertEqual(get_meta_with_checks(32, 32, 32, warn_count=1), default_meta)

        # Test (no)warn_once when tuned parameters are available
        update_meta(32, 32, 48, (2, 8, 5, 6))
        expected_meta = dict(GROUP_SIZE_ROW=2, SPLIT_N=8, num_stages=5, num_warps=6)
        self.assertEqual(get_meta_with_checks(32, 32, 48, warn_count=0), expected_meta)

        # Test non-existing tuned parameters with non-default sparsity
        # while for default sparsity 0.5 the parameters are available
        self.assertEqual(get_meta_with_checks(32, 32, 48, warn_count=0, sparsity=0.6), expected_meta)

        # Test non-existing tuned parameters while there exists
        # parameters with consistent N // SPLIT_N ratio:
        self.assertEqual(get_meta_with_checks(32, 32, 72, warn_count=0),
                         dict(GROUP_SIZE_ROW=2, SPLIT_N=12, num_stages=5, num_warps=6))
        # ... or not:
        self.assertEqual(get_meta_with_checks(32, 32, 64, warn_count=1),
                         dict(GROUP_SIZE_ROW=4, SPLIT_N=4, num_stages=1, num_warps=4))


# e.g., TestSparseCSRCPU and TestSparseCSRCUDA
instantiate_device_type_tests(TestSparseCSR, globals())
instantiate_device_type_tests(TestSparseCompressed, globals())
instantiate_device_type_tests(TestSparseCompressedTritonKernels, globals())

if __name__ == '__main__':
    run_tests()