embedding.h

#pragma once

#include <torch/nn/cloneable.h>
#include <torch/nn/pimpl.h>
#include <torch/types.h>

#include <cstddef>
#include <vector>

namespace torch {
namespace nn {

struct TORCH_API EmbeddingOptions {
  EmbeddingOptions(int64_t count, int64_t dimension);
  TORCH_ARG(int64_t, count);
  TORCH_ARG(int64_t, dimension);
};

class TORCH_API EmbeddingImpl : public torch::nn::Cloneable<EmbeddingImpl> {
 public:
  EmbeddingImpl(int64_t count, int64_t dimension)
      : EmbeddingImpl(EmbeddingOptions(count, dimension)) {}
  explicit EmbeddingImpl(EmbeddingOptions options);

  void reset() override;

  void pretty_print(std::ostream& stream) const override;

  Tensor forward(const Tensor& indices);

  EmbeddingOptions options;

  Tensor weight;
};

TORCH_MODULE(Embedding);

} // namespace nn
} // namespace torch