code: def _validate_csr_hlo(data, indices, indptr, shape):
      data_type = ir.RankedTensorType(data.type)
      indices_type = ir.RankedTensorType(indices.type)
      indptr_type = ir.RankedTensorType(indptr.type)
  
      (nnz,) = data_type.shape
      assert indices_type.shape == [nnz]
      assert indptr_type.element_type == indices_type.element_type
      assert indptr_type.shape == [shape[0] + 1]
      return data_type.element_type, indices_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 1 for calling a global function with example usage: _validate_csr_hlo(data, indices, indptr, shape) and returns: data_type,  indices_type,  nnz