def batched(iterable, n):
  "Batch data into tuples of length n. The last batch may be shorter."
  # batched('ABCDEFG', 3) --> ABC DEF G
  if n < 1:
    raise ValueError('n must be at least one')
  it = iter(iterable)
  while batch := list(itertools.islice(it, n)):
    yield batch
Batch data into tuples of length n. The last batch may be shorter.
slora/models/peft/triton_kernel/tests/benchmark_utils.py