class_docstr: A state in the search process. It consists of the current loop structure
  and a list of transformation steps used to construct it.
  Each State corresponds to a specific schedule for its ComputeDAG.
  Parameters
  ----------
  state_object : StateObject
      The StateObject corresponding to C++ internal State object.
  dag : ComputeDAG
      The original ComputeDAG of this State.
  Notes
  -----
  This is a wrapper class of StateObject to deal with copy-on-write property
class_docstr_tok: ["A","state","in","the","search","process",".","It","consists","of","the","current","loop","structure","and","a","list","of","transformation","steps","used","to","construct","it",".","Each","State","corresponds","to","a","specific","schedule","for","its","ComputeDAG",".","Parameters","--","--","--","--","--","state_object",":","StateObject","The","StateObject","corresponding","to","C++","internal","State","object",".","dag",":","ComputeDAG","The","original","ComputeDAG","of","this","State",".","Notes","--","--","-","This","is","a","wrapper","class","of","StateObject","to","deal","with","copy-on-write","property"]
docstr: Schedule primitive corresponding to `te.Schedule.cache_write`.
  See also the `te.Schedule` for more details.
  Parameters
  ----------
  stage : Union[int, Operation, Tensor]
      The Stage to be cache_write, which can be specified by the integer index, Operation,
      or output tensor of the stage.
  scope_name : str
      The scope name of the newly added compute stage.
  Returns
  -------
  new_stage_op : Operator
      The Operator of the new added stage.
  Notes
  -----
  Cache write step will insert an extra stage to the original ComputeDAG (in the front of the
  target stage).
  This step will cache write all output tensors of the target stage.
docstr_tok: ["Schedule","primitive","corresponding","to","`","te.Schedule.cache_write","`",".","See","also","the","`","te.Schedule","`","for","more","details",".","Parameters","--","--","--","--","--","stage",":","Union","[","int",",","Operation",",","Tensor","]","The","Stage","to","be","cache_write",",","which","can","be","specified","by","the","integer","index",",","Operation",",","or","output","tensor","of","the","stage",".","scope_name",":","str","The","scope","name","of","the","newly","added","compute","stage",".","Returns","--","--","--","-","new_stage_op",":","Operator","The","Operator","of","the","new","added","stage",".","Notes","--","--","-","Cache","write","step","will","insert","an","extra","stage","to","the","original","ComputeDAG","(","in","the","front","of","the","target","stage",")",".","This","step","will","cache","write","all","output","tensors","of","the","target","stage","."]
code: def cache_write(self, stage, scope_name):
      """Schedule primitive corresponding to `te.Schedule.cache_write`.
      See also the `te.Schedule` for more details.
      Parameters
      ----------
      stage : Union[int, Operation, Tensor]
          The Stage to be cache_write, which can be specified by the integer index, Operation,
          or output tensor of the stage.
      scope_name : str
          The scope name of the newly added compute stage.
      Returns
      -------
      new_stage_op : Operator
          The Operator of the new added stage.
      Notes
      -----
      Cache write step will insert an extra stage to the original ComputeDAG (in the front of the
      target stage).
      This step will cache write all output tensors of the target stage.
      """
      self.state_object, new_stage_id = _ffi_api.StateCacheWrite(
          self.state_object,
          self._resolve_stage_id(stage),
          scope_name,
          self.compute_dag,
      )
      # Add a new stage will change all ops behind the added stage. But we still want to keep the
      # original ops map, apply stage id offset to stage_id_map to make them work.
      self._apply_stage_id_offset(int(new_stage_id))
      self._update_stage_id_map()
      return self.stages[int(new_stage_id)].op
  
imports: import tvm._ffi
  from tvm.te.tensor import Operation, Tensor
  from tvm.runtime import Object
  from .None import _ffi_api
label_desc: Use image node_id 20 for calling the State obj's underlying member method code with example usage: obj.cache_write(stage, scope_name) and returns: self