src/modules/data/dataset.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/dataset.h"

#include <cstring>
#include <vector>

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

namespace tenzr::data {

  TensorDataset::TensorDataset(std::vector<Tensor> tensors) : tensors_(std::move(tensors)) {
    TENZR_CHECK(!tensors_.empty(), "TensorDataset cannot be initialized with empty tensors list");
    TENZR_CHECK(tensors_[0].shape().Rank() >= 1, "TensorDataset inputs must have rank >= 1");

    size_ = static_cast<size_t>(tensors_[0].shape().Dim(0));
    for (size_t i = 1; i < tensors_.size(); ++i) {
      TENZR_CHECK(tensors_[i].shape().Rank() >= 1, "TensorDataset inputs must have rank >= 1");
      TENZR_CHECK(static_cast<size_t>(tensors_[i].shape().Dim(0)) == size_,
                  "Dimension mismatch in TensorDataset: tensor 0 has dim(0)=" + std::to_string(size_) + " but tensor " + std::to_string(i) +
                      " has dim(0)=" + std::to_string(tensors_[i].shape().Dim(0)));
    }
  }

  DataItem TensorDataset::Get(size_t index) const {
    TENZR_CHECK(index < size_, "Index " + std::to_string(index) + " out of bounds for dataset of size " + std::to_string(size_));

    DataItem item;
    item.reserve(tensors_.size());

    for (const auto& t : tensors_) {
      std::vector<int64_t> item_dims;
      for (size_t d = 1; d < t.shape().Rank(); ++d) {
        item_dims.push_back(t.shape().Dim(d));
      }
      if (item_dims.empty()) {
        item_dims.push_back(1);
      }
      Shape item_shape(item_dims);

      size_t item_elems = static_cast<size_t>(item_shape.NumElements());
      size_t item_bytes = item_elems * SizeOf(t.Dtype());

      auto slice_tensor = Tensor::Parameter(item_shape, t.Dtype(), "item");

      if (t.HasData()) {
        auto buf = MemoryPool::Instance().Allocate(item_bytes);
        const uint8_t* src_ptr = static_cast<const uint8_t*>(t.Data()->Data()) + (index * item_bytes);
        std::memcpy(buf->Data(), src_ptr, item_bytes);
        slice_tensor.BindData(buf);
      }

      item.push_back(std::move(slice_tensor));
    }

    return item;
  }

  ShardedDataset::ShardedDataset(std::shared_ptr<Dataset> base_dataset, size_t num_shards, size_t shard_index)
      : base_dataset_(std::move(base_dataset)), num_shards_(num_shards), shard_index_(shard_index) {
    TENZR_CHECK(base_dataset_ != nullptr, "Base dataset cannot be null in ShardedDataset");
    TENZR_CHECK(num_shards_ > 0, "num_shards must be > 0");
    TENZR_CHECK(shard_index_ < num_shards_, "shard_index must be < num_shards");

    size_t total = base_dataset_->Size();
    for (size_t i = shard_index_; i < total; i += num_shards_) {
      sharded_indices_.push_back(i);
    }
  }

}  // namespace tenzr::data