embedding_layer.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/layers/layer.h"
#include "opennn/network/operators/dropout_operator.h"
#include "opennn/network/operators/embedding_lookup_operator.h"
 
namespace opennn
{
 
class Embedding final : public Layer
{
public:
 
    Embedding(const Shape& = {0, 0},
              Index = 0,
              const string& = "embedding_layer");
 
    bool accepts_input_rank(Index rank) const override { return is_one_of(rank, 1); }
    bool allows_bf16_input_cast(size_t) const noexcept override { return false; }
 
    Shape get_input_shape() const noexcept override { return {sequence_length}; }
 
    void apply_input_shape(const Shape& new_input_shape) override
    {
        set(vocabulary_size, new_input_shape.dim_or_zero(0), embedding_dimension, label);
    }
    void set_output_shape(const Shape& shape) override
    {
        set(vocabulary_size, shape.dim_or_zero(0), shape.dim_or_zero(1), label);
    }
    Shape get_output_shape() const override;
 
    Index get_vocabulary_size() const { return vocabulary_size; }
    Index get_sequence_length() const { return sequence_length; }
    Index get_embedding_dimension() const { return embedding_dimension; }
 
    vector<TensorSpec> get_forward_specs(Index) const override;
    vector<TensorSpec> get_backward_specs(Index) const override { return {}; }
    ForwardSlotKind get_forward_slot_kind(size_t slot) const override
    {
        return slot == DropoutMask
            ? ForwardSlotKind::TrainingOnly
            : ForwardSlotKind::Pooled;
    }
 
    void set(Index = 0,
             Index = 0,
             Index = 0,
             const string& = "embedding_layer");
 
    void set_scale_embedding(bool enabled) { embedding_lookup.scale_embedding = enabled; }
    void set_add_positional_encoding(bool enabled) { embedding_lookup.add_positional_encoding = enabled; }
 
    void set_learned_positional(bool enabled)
    {
        embedding_lookup.add_positional_encoding = enabled;
        embedding_lookup.positional_trainable    = enabled;
    }
    bool get_learned_positional() const { return embedding_lookup.positional_trainable; }
 
    void set_export_valid_lengths(bool enabled) { embedding_lookup.export_valid_lengths = enabled; }
    bool get_export_valid_lengths() const { return embedding_lookup.export_valid_lengths; }
    void set_weights_follow_compute_dtype(bool enabled) { embedding_lookup.weights_follow_compute_dtype = enabled; }
    void set_dropout_rate(float rate) { dropout.set_rate(rate); }
 
    void on_compute_dtype_changed() override { embedding_lookup.compute_dtype = get_compute_dtype(); }
 
    void read_JSON_body(const Json*) override;
    void write_JSON_body(JsonWriter&) const override;
 
private:
 
    Index vocabulary_size = 0;
    Index sequence_length = 0;
    Index embedding_dimension = 0;
 
    EmbeddingLookupOperator embedding_lookup;
    DropoutOperator         dropout;
 
    enum Forward {Input, DropoutMask, Output};
};
 
}