code: def __init__(
      self,
      fn,
      arg_names,
      configs,
      key,
      reset_to_zero,
      prune_configs_by: Dict = None,
      nearest_power_of_two: bool = False,
  ):
      if not configs:
          self.configs = [triton.Config({}, num_warps=4, num_stages=2)]
      else:
          self.configs = configs
      self.key_idx = [arg_names.index(k) for k in key]
      self.nearest_power_of_two = nearest_power_of_two
      self.cache = {}
      # hook to reset all required tensor to zeros before relaunching a kernel
      self.hook = lambda args: 0
      if reset_to_zero is not None:
          self.reset_idx = [arg_names.index(k) for k in reset_to_zero]
  
          def _hook(args):
              for i in self.reset_idx:
                  args[i].zero_()
  
          self.hook = _hook
      self.arg_names = arg_names
      # prune configs
      if prune_configs_by:
          perf_model, top_k = (
              prune_configs_by["perf_model"],
              prune_configs_by["top_k"],
          )
          if "early_config_prune" in prune_configs_by:
              early_config_prune = prune_configs_by[
                  "early_config_prune"
              ]
      else:
          perf_model, top_k, early_config_prune = None, None, None
      self.perf_model, self.configs_top_k = perf_model, top_k
      self.early_config_prune = early_config_prune
      self.fn = fn
  
file_path: colossalai/colossalai/kernel/triton/custom_autotune.py
imports: import builtins
  import math
  import time
  from typing import Dict
  import triton
label_desc: Use image node_id 1 to create a new CustomizedTritonAutoTuner object from inherited base classes: triton with example: obj = CustomizedTritonAutoTuner(fn, arg_names, configs, key, reset_to_zero, prune_configs_by, nearest_power_of_two)