include/tenzr/modules/data/dataset.h¶

// 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_DATASET_H_
#define TENZR_MODULES_DATA_DATASET_H_

#include <cstddef>
#include <memory>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/frontend/tensor.h"

namespace tenzr::data {

  using DataItem = std::vector<Tensor>;

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

    virtual size_t Size() const = 0;
    virtual DataItem Get(size_t index) const = 0;
  };

  class TensorDataset : public Dataset {
   public:
    explicit TensorDataset(std::vector<Tensor> tensors);

    size_t Size() const override { return size_; }

    DataItem Get(size_t index) const override;

    const std::vector<Tensor>& Tensors() const { return tensors_; }

   private:
    std::vector<Tensor> tensors_;
    size_t size_{0};
  };

  class ShardedDataset : public Dataset {
   public:
    ShardedDataset(std::shared_ptr<Dataset> base_dataset, size_t num_shards, size_t shard_index);

    size_t Size() const override { return sharded_indices_.size(); }

    DataItem Get(size_t index) const override {
      TENZR_CHECK(index < sharded_indices_.size(), "Index out of bounds in ShardedDataset");
      return base_dataset_->Get(sharded_indices_[index]);
    }

    size_t NumShards() const { return num_shards_; }
    size_t ShardIndex() const { return shard_index_; }

   private:
    std::shared_ptr<Dataset> base_dataset_;
    size_t num_shards_;
    size_t shard_index_;
    std::vector<size_t> sharded_indices_;
  };

}  // namespace tenzr::data

#endif  // TENZR_MODULES_DATA_DATASET_H_