code: def test_sample_categorical():
      """Test sample categorical sampling function"""
      n = 1000
      sch = tir.Schedule(elementwise, seed=42, debug_mask="all")
      counter = defaultdict(int)
      candidates = [5, 2, 7, 1]
      probs = [0.15, 0.55, 0.05, 0.25]
      for _ in range(n):
          v = sch.get(sch.sample_categorical(candidates, probs))
          counter[v] += 1
      for i, prob in enumerate(probs):
          assert (
              (prob - 0.07) * n
              <= counter[candidates[i]]
              <= (prob + 0.07) * n
          )
      verify_trace_roundtrip(sch, mod=elementwise)
  
file_path: tvm/tests/python/tir-schedule/test_tir_schedule_sampling.py
imports: from collections import defaultdict
  import sys
  import numpy
  import pytest
  import tvm.testing
  from tvm import tir
  from tvm.script import tir
  from tvm.tir.schedule.testing import verify_trace_roundtrip
label_desc: Use image node_id 3 for calling a global function with example usage: test_sample_categorical() without return types