code: def _match_bucket_keys(bucket_keys, seq_lengths):
      bucket_key_npy = np.array(bucket_keys, dtype=np.int64)
      bucket_sample_ids = [list() for _ in range(len(bucket_keys))]
      batch_size = 10000
      bucket_key_npy = bucket_key_npy.reshape(
          (1,) + bucket_key_npy.shape
      )
      for begin in range(0, len(seq_lengths), batch_size):
          end = min(begin + batch_size, len(seq_lengths))
          diff = bucket_key_npy - np.expand_dims(
              seq_lengths[begin:end], axis=1
          )
          if diff.ndim == 3:
              is_valid_bucket = np.prod(diff >= 0, axis=2)
              pad_val = np.sum(diff, axis=2)
          else:
              is_valid_bucket = diff >= 0
              pad_val = diff
          seq_ids_not_found = np.nonzero(
              is_valid_bucket.sum(axis=1) == 0
          )[0].tolist()
          masked_pad_val = np.ma.array(
              pad_val, mask=1 - is_valid_bucket
          )
          batch_bucket_id = masked_pad_val.argmin(axis=1).tolist()
          if len(seq_ids_not_found) > 0:
              raise ValueError(
                  "Find elements in seq_lengths that cannot fit in the "
                  "given buckets, seq_length={}, bucket_keys={}. "
                  "You must increase the bucket size."
                  % (seq_lengths[seq_ids_not_found], bucket_keys)
              )
          for i, bucket_id in enumerate(batch_bucket_id):
              bucket_sample_ids[bucket_id].append(i + begin)
      return bucket_sample_ids
  
imports: import math
  import random
  import warnings
  import numpy
  import abc
  import logging
  from typing import Union, Sequence, Optional, List
  from ..base import INT_TYPES
label_desc: Use image node_id 1 for calling a global function with example usage: _match_bucket_keys(bucket_keys, seq_lengths) and returns: bucket_sample_ids