Skip to content

include/tenzr/modules/data/sampler.h

Namespaces

Name
tenzr
tenzr::data

Classes

Name
class tenzr::data::Sampler
Abstract source of index orders over a dataset.
class tenzr::data::SequentialSampler
Sampler that yields indices in ascending order.
class tenzr::data::RandomSampler
Sampler that yields a shuffled permutation of the dataset indices.
class tenzr::data::BatchSampler
Groups indices produced by a Sampler into fixed-size batches.

Source code

// 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.

#ifndef TENZR_MODULES_DATA_SAMPLER_H_
#define TENZR_MODULES_DATA_SAMPLER_H_

#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <vector>

namespace tenzr::data {

  class Sampler {
   public:
    virtual ~Sampler() = default;

    virtual std::vector<size_t> Indices() const = 0;
    virtual size_t Size() const = 0;
  };

  class SequentialSampler : public Sampler {
   public:
    explicit SequentialSampler(size_t dataset_size);

    std::vector<size_t> Indices() const override;
    size_t Size() const override { return size_; }

   private:
    size_t size_;
  };

  class RandomSampler : public Sampler {
   public:
    explicit RandomSampler(size_t dataset_size, std::optional<uint64_t> seed = std::nullopt);

    std::vector<size_t> Indices() const override;
    size_t Size() const override { return size_; }

   private:
    size_t size_;
    std::optional<uint64_t> seed_;
  };

  class BatchSampler {
   public:
    BatchSampler(std::shared_ptr<Sampler> sampler, size_t batch_size, bool drop_last = false);

    std::vector<std::vector<size_t>> Batches() const;

    size_t NumBatches() const;
    size_t BatchSize() const { return batch_size_; }
    bool DropLast() const { return drop_last_; }

   private:
    std::shared_ptr<Sampler> sampler_;
    size_t batch_size_;
    bool drop_last_;
  };

}  // namespace tenzr::data

#endif  // TENZR_MODULES_DATA_SAMPLER_H_