Method dynamic_batch_

mlx/data/buffer/DynamicBatch.cpp:68–207

      skipped_samples_(std::get<2>(buffer_with_sizes)) {};

std::tuple<std::shared_ptr<Buffer>, std::vector<int64_t>, std::vector<int64_t>>
DynamicBatch::dynamic_batch_(
    const std::shared_ptr<Buffer>& buffer,
    const std::shared_ptr<Buffer>& ref_size_buffer,
    const std::string& key,
    int64_t min_data_size,
    int64_t max_data_size,
    const std::unordered_map<std::string, int>& batch_dims,
    bool drop_outliers) {
  int64_t n = buffer->size();

  if (ref_size_buffer && (ref_size_buffer->size() != n)) {
    throw std::runtime_error(
        "DynamicBatch: buffer and reference size buffer do not match in size");
  }

  // get sample shapes
  std::vector<std::vector<int64_t>> sample_shapes(n);
  for (int64_t i = 0; i < n; i++) {
    if (ref_size_buffer) {
      auto sample = ref_size_buffer->get(i);
      auto array = mlx::data::sample::check_key(sample, key, ArrayType::Int64);
      sample_shapes[i].insert(
          sample_shapes[i].begin(),
          array->data<int64_t>(),
          array->data<int64_t>() + array->size());
    } else {
      auto sample = buffer->get(i);
      auto array = mlx::data::sample::check_key(sample, key, ArrayType::Any);
      sample_shapes[i] = array->shape();
    }
  }

  // get sample sizes
  std::vector<int64_t> sample_sizes(n, 0);
  for (int64_t i = 0; i < n; i++) {
    auto& shape = sample_shapes[i];
    if (shape.size() > 0) {
      int64_t size = shape[0];
      for (int i = 1; i < shape.size(); i++) {
        size *= shape[i];
      }
      sample_sizes[i] = size;
    }
  }

  // sort (if we do not batch everything)
  std::vector<int64_t> sorted_samples(sample_sizes.size());
  std::iota(sorted_samples.begin(), sorted_samples.end(), 0);
  if (max_data_size > 0) {
    std::stable_sort(
        sorted_samples.begin(),
        sorted_samples.end(),
        [&sample_sizes](size_t i1, size_t i2) {
          return sample_sizes[i1] < sample_sizes[i2];
        });
  }