embedding_lookup_operator.h

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-only
// Copyright (C) 2026 Artificial Intelligence Techniques, S.L. and Roberto López González
 
#pragma once
 
#include "opennn/network/operators/operator.h"
 
namespace opennn
{
 
void embedding_lookup_forward(const TensorView&, const TensorView&,
                              const TensorView&, TensorView&,
                              Index, Index, Index,
                              bool scale_embedding, bool add_positional_encoding,
                              const TensorView& weight_scale = {});
void embedding_lookup_backward(const TensorView&, const TensorView&,
                               const TensorView&, const TensorView&,
                               Index, Index, Index,
                               bool scale_embedding);
void compute_token_valid_lengths(const TensorView&, Index, vector<Index>&);
 
struct EmbeddingLookupOperator : Operator
{
    Index vocabulary_size     = 0;
    Index sequence_length     = 0;
    Index embedding_dimension = 0;
 
    bool scale_embedding         = false;
    bool add_positional_encoding = false;
    bool  positional_trainable    = false;
    bool  export_valid_lengths    = false;
 
    bool weights_follow_compute_dtype = false;
 
    TensorView weights;
    TensorView weight_scale;
    TensorView positional_encoding;
 
    TensorView weight_gradient;
    TensorView positional_gradient;
 
    void set(Index, Index, Index);
 
    vector<TensorSpec> parameter_specs() const override;
    vector<SlotQuantization> parameter_quantization() const override;
    vector<TensorSpec> state_specs()     const override;
    vector<ParameterSlot> parameter_slots() override;
    void link_states    (span<const TensorView>) override;
    void initialize_states() override;
    void link_parameter_scales(span<const TensorView>) override;
 
    void set_parameters_random() override;
    void set_parameters_glorot() override;
 
    void init_positional_encoding();
    void init_trainable_positional();
 
    void forward_propagate(ForwardPropagation&, size_t, ForwardPropagationMode) override;
    void back_propagate(ForwardPropagation&, BackPropagation&, size_t) const override;
 
    void load_state_from_JSON(const Json*) override;
};
 
}