batch_embedding_lookup_expr.hpp

//=======================================================================
// Copyright (c) 2014-2023 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

#include "etl/expr/base_temporary_expr.hpp"

namespace etl {

template <etl_2d A, etl_2d B>
struct batch_embedding_lookup_expr : base_temporary_expr_bin<batch_embedding_lookup_expr<A, B>, A, B> {
    using value_type = value_t<A>;                               
    using this_type  = batch_embedding_lookup_expr<A, B>;        
    using base_type  = base_temporary_expr_bin<this_type, A, B>; 
    using sub_traits = decay_traits<A>;                          

    static constexpr auto storage_order = sub_traits::storage_order; 

    static constexpr bool gpu_computable = false;

    explicit batch_embedding_lookup_expr(A a, B b) : base_type(a, b) {
        //Nothing else to init
    }

    template <etl_3d C>
    static void check([[maybe_unused]] const A& a, [[maybe_unused]] const B& b, [[maybe_unused]] const C& c) {
        if constexpr (all_fast<A, B, C>) {
            static_assert(etl::dim<0, A>() == etl::dim<0, C>(), "Invalid dimensions for batch_embedding_lookup");
            static_assert(etl::dim<1, A>() == etl::dim<1, C>(), "Invalid dimensions for batch_embedding_lookup");
            static_assert(etl::dim<1, B>() == etl::dim<2, C>(), "Invalid dimensions for batch_embedding_lookup");
        } else {
            cpp_assert(etl::dim<0>(a) == etl::dim<0>(c), "Invalid dimensions for batch_embedding_lookup");
            cpp_assert(etl::dim<1>(a) == etl::dim<1>(c), "Invalid dimensions for batch_embedding_lookup");
            cpp_assert(etl::dim<1>(b) == etl::dim<2>(c), "Invalid dimensions for batch_embedding_lookup");
        }
    }

    // Assignment functions

    template <etl_3d L>
    void assign_to(L&& lhs) const {
        inc_counter("temp:assign");

        auto& a = this->a();
        auto& b = this->b();

        check(a, b, lhs);

        standard_evaluator::pre_assign_rhs(a);
        standard_evaluator::pre_assign_rhs(b);

        const auto BB = etl::dim<0>(a);
        const auto I  = etl::dim<1>(a);

        for (size_t bb = 0; bb < BB; ++bb) {
            for (size_t i = 0; i < I; ++i) {
                lhs(bb)(i) = b(a(bb, i));
            }
        }
    }

    template <etl_3d L>
    void assign_add_to(L&& lhs) const {
        auto& a = this->a();
        auto& b = this->b();

        check(a, b, lhs);

        standard_evaluator::pre_assign_rhs(a);
        standard_evaluator::pre_assign_rhs(b);

        const auto BB = etl::dim<0>(a);
        const auto I  = etl::dim<1>(a);

        for (size_t bb = 0; bb < BB; ++bb) {
            for (size_t i = 0; i < I; ++i) {
                lhs(bb)(i) += b(a(bb, i));
            }
        }
    }

    template <etl_3d L>
    void assign_sub_to(L&& lhs) const {
        auto& a = this->a();
        auto& b = this->b();

        check(a, b, lhs);

        standard_evaluator::pre_assign_rhs(a);
        standard_evaluator::pre_assign_rhs(b);

        const auto BB = etl::dim<0>(a);
        const auto I  = etl::dim<1>(a);

        for (size_t bb = 0; bb < BB; ++bb) {
            for (size_t i = 0; i < I; ++i) {
                lhs(bb)(i) -= b(a(bb, i));
            }
        }
    }

    template <etl_3d L>
    void assign_mul_to(L&& lhs) const {
        auto& a = this->a();
        auto& b = this->b();

        check(a, b, lhs);

        standard_evaluator::pre_assign_rhs(a);
        standard_evaluator::pre_assign_rhs(b);

        const auto BB = etl::dim<0>(a);
        const auto I  = etl::dim<1>(a);

        for (size_t bb = 0; bb < BB; ++bb) {
            for (size_t i = 0; i < I; ++i) {
                lhs(bb)(i) *= b(a(bb, i));
            }
        }
    }

    template <etl_3d L>
    void assign_div_to(L&& lhs) const {
        auto& a = this->a();
        auto& b = this->b();

        check(a, b, lhs);

        standard_evaluator::pre_assign_rhs(a);
        standard_evaluator::pre_assign_rhs(b);

        const auto BB = etl::dim<0>(a);
        const auto I  = etl::dim<1>(a);

        for (size_t bb = 0; bb < BB; ++bb) {
            for (size_t i = 0; i < I; ++i) {
                lhs(bb)(i) /= bb(a(b, i));
            }
        }
    }

    template <etl_3d L>
    void assign_mod_to(L&& lhs) const {
        auto& a = this->a();
        auto& b = this->b();

        check(a, b, lhs);

        standard_evaluator::pre_assign_rhs(a);
        standard_evaluator::pre_assign_rhs(b);

        const auto BB = etl::dim<0>(a);
        const auto I  = etl::dim<1>(a);

        for (size_t bb = 0; bb < BB; ++bb) {
            for (size_t i = 0; i < I; ++i) {
                lhs(bb)(i) %= b(a(bb, i));
            }
        }
    }

    friend std::ostream& operator<<(std::ostream& os, const batch_embedding_lookup_expr& expr) {
        return os << "embedding_lookup(" << expr._a << ", " << expr._b << ")";
    }
};

template <etl_expr A, etl_expr B>
struct etl_traits<etl::batch_embedding_lookup_expr<A, B>> {
    using expr_t     = etl::batch_embedding_lookup_expr<A, B>; 
    using sub_expr_t = std::decay_t<A>;                        
    using sub_traits = etl_traits<sub_expr_t>;                 
    using value_type = value_t<A>;                             

    static constexpr bool is_etl         = true;                      
    static constexpr bool is_transformer = false;                     
    static constexpr bool is_view        = false;                     
    static constexpr bool is_magic_view  = false;                     
    static constexpr bool is_fast        = sub_traits::is_fast;       
    static constexpr bool is_linear      = false;                     
    static constexpr bool is_thread_safe = true;                      
    static constexpr bool is_value       = false;                     
    static constexpr bool is_direct      = true;                      
    static constexpr bool is_generator   = false;                     
    static constexpr bool is_padded      = false;                     
    static constexpr bool is_aligned     = true;                      
    static constexpr bool is_temporary   = true;                      
    static constexpr order storage_order = sub_traits::storage_order; 
    static constexpr bool gpu_computable = false;                     

    template <vector_mode_t V>
    static constexpr bool vectorizable = true;

    template <size_t DD>
    static constexpr size_t dim() requires(DD < 3) {
        if (DD == 0) {
            return decay_traits<A>::template dim<0>();
        } else if (DD == 1) {
            return decay_traits<A>::template dim<1>();
        } else {
            return decay_traits<B>::template dim<1>();
        }
    }

    static size_t dim(const expr_t& e, [[maybe_unused]] size_t d) {
        cpp_assert(d < 3, "Invalid dimensions access");

        if (d == 0) {
            return etl::dim<0>(e._a);
        } else if (d == 1) {
            return etl::dim<1>(e._a);
        } else {
            return etl::dim<1>(e._b);
        }
    }

    static size_t size(const expr_t& e) {
        return etl::dim<0>(e._a) * etl::dim<1>(e._a) * etl::dim<1>(e._b);
    }

    static constexpr size_t size() {
        return decay_traits<A>::template dim<0>() * decay_traits<A>::template dim<1>() * decay_traits<B>::template dim<1>();
    }

    static constexpr size_t dimensions() {
        return 3;
    }

    static constexpr int complexity() noexcept {
        return 4;
    }
};

template <etl_2d I, etl_2d V>
batch_embedding_lookup_expr<detail::build_type<I>, detail::build_type<V>> batch_embedding_lookup(const I& value, const V& vocab) {
    return batch_embedding_lookup_expr<detail::build_type<I>, detail::build_type<V>>{value, vocab};
}

} //end of namespace etl