code: def test_spmm_jit(reduce):
      @torch.jit.script
      def jit_torch_sparse(
          src: SparseTensor, other: Tensor, reduce: str
      ) -> Tensor:
          return spmm(src, other, reduce=reduce)
  
      @torch.jit.script
      def jit_torch(src: Tensor, other: Tensor, reduce: str) -> Tensor:
          return spmm(src, other, reduce=reduce)
  
      src = torch.randn(5, 4)
      other = torch.randn(4, 8)
  
      out1 = src @ other
      out2 = jit_torch(src.to_sparse_csr(), other, reduce)
      assert out1.size() == (5, 8)
      if reduce == "sum":
          assert torch.allclose(out1, out2, atol=1e-6)
      if torch_geometric.typing.WITH_TORCH_SPARSE:
          out3 = jit_torch_sparse(
              SparseTensor.from_dense(src), other, reduce
          )
          assert torch.allclose(out2, out3, atol=1e-6)
  
file_path: pytorch_geometric/test/utils/test_spmm.py
imports: import itertools
  import warnings
  import pytest
  import torch
  from torch import Tensor
  import torch_geometric.typing
  from torch_geometric.profile import benchmark
  from torch_geometric.testing import withCUDA, withPackage
  from torch_geometric.typing import SparseTensor
  from torch_geometric.utils import spmm, to_torch_coo_tensor
label_desc: Use image node_id 4 for calling a global function with example usage: test_spmm_jit(reduce) and returns: spmm, spmm