src/modules/nn/embedding.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/nn/embedding.h"

#include <cstring>
#include <vector>

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

namespace tenzr::nn {

  Embedding::Embedding(int64_t num_embeddings, int64_t embedding_dim, DType dtype, std::optional<uint64_t> seed)
      : num_embeddings_(num_embeddings), embedding_dim_(embedding_dim) {
    TENZR_CHECK(num_embeddings > 0, "num_embeddings must be positive");
    TENZR_CHECK(embedding_dim > 0, "embedding_dim must be positive");

    if (seed.has_value()) {
      random::Generator gen(*seed);
      weight_ = RegisterParameter("weight", random::Normal({num_embeddings_, embedding_dim_}, 0.0f, 1.0f, dtype, gen));
    } else {
      weight_ = RegisterParameter("weight", random::Normal({num_embeddings_, embedding_dim_}, 0.0f, 1.0f, dtype, random::DefaultGenerator()));
    }
  }

  Tensor Embedding::Forward(const Tensor& indices) {
    TENZR_CHECK(indices.IsDefined(), "Indices tensor is uninitialized");

    std::vector<int64_t> out_dims;
    for (size_t i = 0; i < indices.shape().Rank(); ++i) {
      out_dims.push_back(indices.shape().Dim(i));
    }
    out_dims.push_back(embedding_dim_);
    Shape out_shape(out_dims);

    if (indices.HasData() && weight_.HasData()) {
      int64_t num_indices = indices.shape().NumElements();
      std::vector<int64_t> idx_vals(static_cast<size_t>(num_indices));

      if (indices.Dtype() == DType::kInt32) {
        auto raw = indices.ToVector<int32_t>();
        for (size_t i = 0; i < raw.size(); ++i) {
          idx_vals[i] = raw[i];
        }
      } else if (indices.Dtype() == DType::kInt64) {
        idx_vals = indices.ToVector<int64_t>();
      } else {
        TENZR_THROW("Embedding indices must be Int32 or Int64");
      }

      auto w_vec = weight_.ToVectorAs<double>();
      std::vector<double> out_vec(static_cast<size_t>(out_shape.NumElements()));

      for (int64_t i = 0; i < num_indices; ++i) {
        int64_t idx = idx_vals[static_cast<size_t>(i)];
        TENZR_CHECK(idx >= 0 && idx < num_embeddings_, "Index " + std::to_string(idx) + " out of bounds for embedding size " + std::to_string(num_embeddings_));

        size_t src_offset = static_cast<size_t>(idx * embedding_dim_);
        size_t dst_offset = static_cast<size_t>(i * embedding_dim_);
        std::memcpy(&out_vec[dst_offset], &w_vec[src_offset], static_cast<size_t>(embedding_dim_) * sizeof(double));
      }

      return Tensor::FromVector<double>(out_vec, out_shape).Astype(weight_.Dtype());
    }

    return Tensor::CustomOp("embedding", {indices, weight_});
  }

}  // namespace tenzr::nn