code: def search_partioned_ah(
      searcher,
      dims_per_block,
      aiq_threshold,
      reorder_k,
      partioning_trainsize,
      num_leaves,
      num_leaves_to_search,
  ):
      return (
          searcher.tree(
              num_leaves=num_leaves,
              num_leaves_to_search=num_leaves_to_search,
              training_sample_size=partioning_trainsize,
          )
          .score_ah(
              dims_per_block,
              anisotropic_quantization_threshold=aiq_threshold,
          )
          .reorder(reorder_k)
          .build()
      )
  
file_path: ColossalAI/examples/images/diffusion/scripts/train_searcher.py
imports: import argparse
  import glob
  import os
  import sys
  from multiprocessing import cpu_count
  import numpy
  import scann
  from ldm.util import parallel_data_prefetch
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: search_partioned_ah(searcher, dims_per_block, aiq_threshold, reorder_k, partioning_trainsize, num_leaves, num_leaves_to_search) and returns: searcher