Skip to content

tenzr::nn::Embedding

Lookup-table module mapping integer indices to dense vectors.

#include <embedding.h>

Inherits from tenzr::nn::Module

Public Functions

Name
Embedding(int64_t num_embeddings, int64_t embedding_dim, DType dtype, std::optional< uint64_t > seed =std::nullopt)
template <typename T =float>
Embedding(int64_t num_embeddings, int64_t embedding_dim, std::optional< uint64_t > seed =std::nullopt)
Construct with the parameter dtype inferred from T.
template <typename T >
Embedding(std::type_identity< T > , int64_t num_embeddings, int64_t embedding_dim, std::optional< uint64_t > seed =std::nullopt)
Construct from an explicit std::type_identity dtype tag.
virtual Tensor Forward(const Tensor & indices) override
Tensor & Weight()
const Tensor & Weight() const
int64_t NumEmbeddings() const
int64_t EmbeddingDim() const
virtual std::string Name() const override
template <typename T =float>
Embedding
Create(int64_t num_embeddings, int64_t embedding_dim, std::optional< uint64_t > seed =std::nullopt)
Static factory mirroring the templated constructor.

Additional inherited members

Public Functions inherited from tenzr::nn::Module

Name
virtual ~Module() =default
Tensor operator()(const Tensor & input)
virtual void Train(bool mode)
void Train()
Switch to training mode.
void Eval()
Switch to evaluation mode.
bool IsTraining() const
Tensor & RegisterParameter(std::string name, Tensor param)
template <typename T ,typename... Args>
std::shared_ptr< T >
RegisterSubmodule(std::string name, Args &&... args)
std::shared_ptr< Module > RegisterSubmodule(std::string name, std::shared_ptr< Module > submod)
std::vector< Tensor > Parameters(bool recurse =true) const
std::unordered_map< std::string, Tensor > NamedParameters(std::string prefix ="", bool recurse =true) const
std::unordered_map< std::string, Tensor > StateDict() const
void LoadStateDict(const std::unordered_map< std::string, Tensor > & state)

Protected Attributes inherited from tenzr::nn::Module

Name
std::unordered_map< std::string, Tensor > parameters_
std::unordered_map< std::string, std::shared_ptr< Module > > submodules_
bool training_

Public Functions Documentation

function Embedding

Embedding(
    int64_t num_embeddings,
    int64_t embedding_dim,
    DType dtype,
    std::optional< uint64_t > seed =std::nullopt
)

Parameters:

  • num_embeddings Size of the embedding table.
  • embedding_dim Width of each embedding vector.
  • dtype Parameter dtype.
  • seed Optional PRNG seed.
  • num_embeddings Table size.
  • embedding_dim Per-entry vector width.
  • dtype Parameter dtype.
  • seed Optional seed for weight initialization.

Construct with an explicit runtime dtype. Create an embedding table initialized from a standard normal.

function Embedding

template <typename T  =float>
inline Embedding(
    int64_t num_embeddings,
    int64_t embedding_dim,
    std::optional< uint64_t > seed =std::nullopt
)

Construct with the parameter dtype inferred from T.

function Embedding

template <typename T >
inline explicit Embedding(
    std::type_identity< T > ,
    int64_t num_embeddings,
    int64_t embedding_dim,
    std::optional< uint64_t > seed =std::nullopt
)

Construct from an explicit std::type_identity dtype tag.

function Forward

virtual Tensor Forward(
    const Tensor & indices
) override

Parameters:

  • indices Integer index tensor.
  • indices Int32/Int64 index tensor.

Return:

  • Tensor of shape (*indices.shape(), embedding_dim).
  • Embedded tensor of shape indices.shape + [embedding_dim].

Reimplements: tenzr::nn::Module::Forward

Gather rows of the embedding table. Gather embedding rows by index, appending embedding_dim to the index shape.

function Weight

inline Tensor & Weight()

Return: Mutable reference to the embedding table.

function Weight

inline const Tensor & Weight() const

Return: Const reference to the embedding table.

function NumEmbeddings

inline int64_t NumEmbeddings() const

Return: Number of rows in the embedding table.

function EmbeddingDim

inline int64_t EmbeddingDim() const

Return: Width of each embedding vector.

function Name

inline virtual std::string Name() const override

Return: "Embedding".

Reimplements: tenzr::nn::Module::Name

function Create

template <typename T  =float>
static inline Embedding Create(
    int64_t num_embeddings,
    int64_t embedding_dim,
    std::optional< uint64_t > seed =std::nullopt
)

Static factory mirroring the templated constructor.