code: def _csr_matvec_hlo(
      platform,
      gpu_sparse,
      data,
      indices,
      indptr,
      x,
      *,
      shape,
      transpose=False,
      compute_dtype=None,
      compute_type=None,
      data_dtype,
      index_dtype,
      x_dtype,
  ):
      """CSR matrix/vector multiply."""
      data_type, index_type, nnz = _validate_csr_hlo(
          data, indices, indptr, shape
      )
      rows, cols = shape
  
      if compute_dtype is None:
          compute_dtype = data_dtype
          compute_type = data_type
  
      buffer_size, opaque = gpu_sparse.build_csr_matvec_descriptor(
          data_dtype,
          x_dtype,
          compute_dtype,
          index_dtype,
          rows,
          cols,
          nnz,
          transpose,
      )
      out_size = cols if transpose else rows
  
      out = custom_call(
          f"{platform}sparse_csr_matvec",
          result_types=[
              ir.RankedTensorType.get([out_size], compute_type),
              ir.RankedTensorType.get(
                  [buffer_size], ir.IntegerType.get_signless(8)
              ),
          ],
          operands=[data, indices, indptr, x],
          backend_config=opaque,
          operand_layouts=[[0]] * 4,
          result_layouts=[[0]] * 2,
      ).results
      return out[0]
  
file_path: jax/jaxlib/gpu_sparse.py
imports: import math
  from functools import partial
  import importlib
  import jaxlib.mlir.ir
  import numpy
  from jaxlib import xla_client
  from .hlo_helpers import custom_call, mk_result_types_and_shapes
label_desc: Use image node_id 5 for calling a global function with example usage: _csr_matvec_hlo(platform, gpu_sparse, data, indices, indptr, x) and returns: out