code: def _validate_coo_hlo(data, row, col):
      data_type = ir.RankedTensorType(data.type)
      row_type = ir.RankedTensorType(row.type)
      col_type = ir.RankedTensorType(col.type)
  
      (nnz,) = data_type.shape
      assert row_type.shape == [nnz]
      assert col_type.element_type == row_type.element_type
      assert col_type.shape == [nnz]
      return data_type.element_type, row_type.element_type, nnz
  
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 2 for calling a global function with example usage: _validate_coo_hlo(data, row, col) and returns: data_type,  row_type,  nnz