code: def _coo_matmat_hlo(
      platform,
      gpu_sparse,
      data,
      row,
      col,
      B,
      *,
      shape,
      transpose=False,
      compute_dtype=None,
      compute_type=None,
      x_dtype,
      data_dtype,
      index_dtype,
  ):
      """COO from dense matrix."""
      data_type, _, nnz = _validate_coo_hlo(data, row, col)
      is_batched_matmat = False
      batch_count = 1
      if len(shape) == 2:
          rows, cols = shape
      elif len(shape) == 3:
          is_batched_matmat = True
          batch_count, rows, cols = shape
          # Redefine nnz as nnz per batch.
          nnz = nnz // batch_count
  
      B_shape = ir.RankedTensorType(B.type).shape
      _, Ccols = B_shape
  
      if compute_dtype is None:
          compute_dtype = data_dtype
          compute_type = data_type
  
      # TODO(tianjianlu): use batch stride to trigger different mode of batch
      # computation. Currently batch_stride = 0 is not allowed because of the issue
      # in cusparse https://github.com/NVIDIA/CUDALibrarySamples/issues/81#issuecomment-1205562643
      # Set batch stride to be the matrix size for now.
      lhs_batch_stride = nnz
      B_rows = rows if transpose else cols
      rhs_batch_stride = B_rows * Ccols
  
      buffer_size, opaque = gpu_sparse.build_coo_matmat_descriptor(
          data_dtype,
          x_dtype,
          compute_dtype,
          index_dtype,
          rows,
          cols,
          Ccols,
          nnz,
          transpose,
          batch_count,
          lhs_batch_stride,
          rhs_batch_stride,
      )
      out_size = cols if transpose else rows
  
      if is_batched_matmat:
          out_shape = [batch_count, out_size, Ccols]
          out_layout = [2, 1, 0]
      else:
          out_shape = [out_size, Ccols]
          out_layout = [1, 0]
  
      out = custom_call(
          f"{platform}sparse_coo_matmat",
          result_types=[
              ir.RankedTensorType.get(out_shape, compute_type),
              ir.RankedTensorType.get(
                  [buffer_size], ir.IntegerType.get_signless(8)
              ),
          ],
          operands=[data, row, col, B],
          backend_config=opaque,
          operand_layouts=[[0], [0], [0], [1, 0]],
          result_layouts=[out_layout, [0]],
      ).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 10 for calling a global function with example usage: _coo_matmat_hlo(platform, gpu_sparse, data, row, col, B) and returns: out