Method get

mlx/data/buffer/Batch.cpp:52–65

}

Sample Batch::get(int64_t idx) const {
  if (idx < 0 || idx >= size_) {
    throw std::runtime_error("Batch: index out of range");
  }
  auto batch_size =
      (batchSize_ ? std::min(batchSize_, op_->size() - idx * batchSize_)
                  : batchSizes_[idx]);
  auto batch_offset = (batchSize_ ? idx * batchSize_ : batchOffsets_[idx]);
  std::vector<Sample> samples(batch_size);
  for (int64_t i = 0; i < batch_size; i++) {
    samples[i] = op_->get(batch_offset + i);
  }
  return core::merge_batch(samples, padValues_, batchDims_);
}

int64_t Batch::size() const {
  return size_;