src/modules/data/sampler.cc¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "tenzr/modules/data/sampler.h"

#include <algorithm>
#include <numeric>
#include <optional>
#include <random>

#include "tenzr/common/error.h"
#include "tenzr/modules/random/generator.h"

namespace tenzr::data {

  SequentialSampler::SequentialSampler(size_t dataset_size) : size_(dataset_size) {}

  std::vector<size_t> SequentialSampler::Indices() const {
    std::vector<size_t> idxs(size_);
    std::iota(idxs.begin(), idxs.end(), 0);
    return idxs;
  }

  RandomSampler::RandomSampler(size_t dataset_size, std::optional<uint64_t> seed) : size_(dataset_size), seed_(seed) {}

  std::vector<size_t> RandomSampler::Indices() const {
    std::vector<size_t> idxs(size_);
    std::iota(idxs.begin(), idxs.end(), 0);
    uint64_t s = seed_.has_value() ? *seed_ : tenzr::random::DefaultGenerator().NextUint64();
    std::mt19937_64 rng(s);
    std::shuffle(idxs.begin(), idxs.end(), rng);
    return idxs;
  }

  BatchSampler::BatchSampler(std::shared_ptr<Sampler> sampler, size_t batch_size, bool drop_last)
      : sampler_(std::move(sampler)), batch_size_(batch_size), drop_last_(drop_last) {
    TENZR_CHECK(sampler_ != nullptr, "Sampler cannot be null in BatchSampler");
    TENZR_CHECK(batch_size_ > 0, "batch_size must be positive");
  }

  size_t BatchSampler::NumBatches() const {
    size_t total = sampler_->Size();
    if (drop_last_) {
      return total / batch_size_;
    }
    return (total + batch_size_ - 1) / batch_size_;
  }

  std::vector<std::vector<size_t>> BatchSampler::Batches() const {
    auto all_indices = sampler_->Indices();
    std::vector<std::vector<size_t>> result;
    result.reserve(NumBatches());

    std::vector<size_t> current_batch;
    current_batch.reserve(batch_size_);

    for (size_t idx : all_indices) {
      current_batch.push_back(idx);
      if (current_batch.size() == batch_size_) {
        result.push_back(std::move(current_batch));
        current_batch.clear();
        current_batch.reserve(batch_size_);
      }
    }

    if (!current_batch.empty() && !drop_last_) {
      result.push_back(std::move(current_batch));
    }

    return result;
  }

}  // namespace tenzr::data