code: def __init__(self, dtype, batches, D, keep_ratio, bw: bool) -> None:
      dtype_str = DTYPE2STR.get(dtype, dtype)
      self.sub_label = (
          f"{dtype_str} D={D} batches={batches} keep={keep_ratio}"
      )
      self.label = "index_select"
  
      indices = []
      sources = []
      for B, seqlen in batches:
          index = [i for i in range(B)]
          random.Random(B).shuffle(index)
          indices.append(
              torch.zeros(
                  index[int(keep_ratio * B)],
                  dtype=torch.int64,
                  device="cuda",
              )
          )
          source_i = torch.randn(
              [B, seqlen * D],
              dtype=dtype,
              device="cuda",
              requires_grad=bw,
          )
          sources.append(source_i)
      self.indices, self.sources = indices, sources
      self.out = torch.Tensor()
  
file_path: xformers/xformers/benchmarks/benchmark_indexing.py
imports: import random
  import torch
  from utils import DTYPE2STR, benchmark_main_helper2, product_dict
  import xformers.ops
label_desc: Use image node_id 1 to create a new IndexSelectBenchmark object  with example: obj = IndexSelectBenchmark(dtype, batches, D, keep_ratio, bw)