code: def argmax_stir_schedule_nhwc(func, in_layout, out_layout):
      """Schedule for nhwc argmax"""
      sch = tir.Schedule(func, debug_mask="all")
      sch.transform_layout("A_red_temp", "A", in_layout)
      sch.transform_layout("A_red", "A_red", out_layout)
      return sch
  
imports: from tvm import tir
  from tvm import topi
  from ..utils import get_layout_transform_fn
label_desc: Use image node_id 2 for calling a global function with example usage: argmax_stir_schedule_nhwc(func, in_layout, out_layout) and returns: sch