code: def test_torch_spmm(device, reduce, transpose, is_undirected):
      if is_undirected:
          kwargs = dict(is_undirected=True)
          adj = EdgeIndex(
              [[0, 1, 1, 2], [1, 0, 2, 1]], device=device, **kwargs
          )
      else:
          adj = EdgeIndex([[0, 1, 1, 2], [2, 0, 1, 2]], device=device)
      adj, perm = adj.sort_by("col" if transpose else "row")
  
      # Basic:
      x = torch.randn(3, 2, device=device)
  
      if (
          not x.is_cuda and torch_geometric.typing.WITH_PT20
      ) or reduce in ["sum", "add"]:
          out = _TorchSPMM.apply(adj, x, None, reduce, transpose)
          exp = _scatter_spmm(adj, x, None, reduce, transpose)
          assert out.allclose(exp)
      else:
          with pytest.raises(AssertionError):
              _TorchSPMM.apply(adj, x, None, reduce, transpose)
  
      # With non-zero values:
      x = torch.randn(3, 1, device=device)
      value = torch.rand(adj.size(1), device=device)
  
      if (
          not x.is_cuda and torch_geometric.typing.WITH_PT20
      ) or reduce in ["sum", "add"]:
          out = _TorchSPMM.apply(adj, x, value, reduce, transpose)
          exp = _scatter_spmm(adj, x, value, reduce, transpose)
          assert out.allclose(exp)
      else:
          with pytest.raises(AssertionError):
              _TorchSPMM.apply(adj, x, value, reduce, transpose)
  
      # Gradients w.r.t. other:
      x1 = torch.randn(3, 1, device=device, requires_grad=True)
      x2 = x1.detach().requires_grad_()
      grad = torch.randn_like(x1)
  
      if reduce in ["sum", "add"]:
          out = _TorchSPMM.apply(adj, x1, None, reduce, transpose)
          out.backward(grad)
          exp = _scatter_spmm(adj, x2, None, reduce, transpose)
          exp.backward(grad)
          assert x1.grad.allclose(x2.grad)
      else:
          with pytest.raises(AssertionError):
              out = _TorchSPMM.apply(adj, x1, None, reduce, transpose)
              out.backward(grad)
  
      # Gradients w.r.t. value:
      x = torch.randn(3, 1, device=device)
      value1 = torch.rand(
          adj.size(1), device=device, requires_grad=True
      )
      grad = torch.randn_like(x)
  
      with pytest.raises((AssertionError, NotImplementedError)):
          out = _TorchSPMM.apply(adj, x, value1, reduce, transpose)
          out.backward(grad)
  
imports: import os.path
  import warnings
  from typing import List, Optional
  import pytest
  import torch
  from torch import Tensor, tensor
  import torch_geometric
  from torch_geometric import EdgeIndex
  from torch_geometric.edge_index import SUPPORTED_DTYPES, ReduceType, SortReturnType, _scatter_spmm, _torch_sparse_spmm, _TorchSPMM, set_tuple_item
  from torch_geometric.profile import benchmark
  from torch_geometric.testing import onlyCUDA, onlyLinux, withCUDA, withoutExtensions, withPackage
  from torch_geometric.typing import SparseTensor
  from torch_geometric.utils import scatter
label_desc: Use image node_id 27 for calling a global function with example usage: test_torch_spmm(device, reduce, transpose, is_undirected) without return types