code: def sparse_dense_apply_func(search_policy, state, stage_id):
      """Describe how to generate the initial sketch for sparse dense"""
      ret = []
      s_0 = auto_scheduler.loop_state.State(
          state, search_policy.search_task.compute_dag
      )
      if (
          s_0.stages[stage_id].op.tag
          == "sparse_dense_sp_rhs_bsrmm_block"
      ):
          return [s_0.state_object, stage_id - 1]
  
      sparse_dense = s_0.stages[stage_id].op
      sparse_dense_block = s_0.stages[stage_id - 1].op
      assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
      assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
  
      # Set the default consumer of compute block
      consumer = sparse_dense
  
      # If sparse dense has a single elementwise consumer
      # We can compute inline the sparse_dense output stage
      consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
          search_policy.search_task, s_0.state_object, stage_id
      )
      if len(consumers) == 1:
          consumer_id = int(consumers.items()[0][0])
          if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
              search_policy.search_task,
              s_0.state_object,
              stage_id,
              consumer_id,
          ):
              consumer = s_0.stages[consumer_id].op
              s_0.compute_inline(sparse_dense)
  
      i, nb_j, j, row_offset, c = s_0[sparse_dense_block].iters
      m, n = s_0[consumer].iters
      i_0, i_1, i_2 = s_0.split(sparse_dense_block, i, [None, None])
      m_0, m_1 = s_0.follow_split(
          consumer, m, len(s_0.transform_steps) - 1, 1
      )
      j_0, j_1 = s_0.split(sparse_dense_block, nb_j, [None])
      n_0, n_1 = s_0.follow_split(
          consumer, n, len(s_0.transform_steps) - 1, 1
      )
      s_0.reorder(
          sparse_dense_block,
          [i_0, j_0, i_1, j_1, row_offset, i_2, j, c],
      )
      s_0.reorder(consumer, [m_0, n_0, m_1, n_1])
      s_0.compute_at(sparse_dense_block, consumer, n_0)
  
      ret.append([s_0.state_object, stage_id - 2])
  
      return ret
  
file_path: tvm/python/tvm/topi/sparse/utils.py
imports: import tvm
  from tvm import relay, auto_scheduler
  from tvm.relay import data_dep_optimization
  from tvm.auto_scheduler import _ffi_api
label_desc: Use image node_id 9 for calling a global function with example usage: sparse_dense_apply_func(search_policy, state, stage_id) and returns: ret, list