reduc_transformers.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

namespace etl {

template <typename T>
struct argmax_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    friend struct etl_traits<argmax_transformer>;

    static constexpr bool gpu_computable = false;

private:
    sub_type sub; 

public:
    explicit argmax_transformer(sub_type expr) : sub(expr) {}

    value_type operator[](size_t i) const {
        return value_type(max_index(sub(i)));
    }

    value_type read_flat(size_t i) const {
        return value_type(max_index(sub(i)));
    }

    template <typename... Sizes>
    value_type operator()(size_t i, Sizes... /*sizes*/) const {
        return value_type(max_index(sub(i)));
    }

    template <typename E>
    bool alias(const E& rhs) const noexcept {
        return sub.alias(rhs);
    }

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(visitor));
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure sub value
        sub.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        sub.ensure_gpu_up_to_date();
    }

    friend std::ostream& operator<<(std::ostream& os, const argmax_transformer& transformer) {
        return os << "argmax(" << transformer.sub << ")";
    }
};

template <typename T>
struct argmin_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    friend struct etl_traits<argmin_transformer>;

    static constexpr bool gpu_computable = false;

private:
    sub_type sub; 

public:
    explicit argmin_transformer(sub_type expr) : sub(expr) {}

    value_type operator[](size_t i) const {
        return value_type(min_index(sub(i)));
    }

    value_type read_flat(size_t i) const {
        return value_type(min_index(sub(i)));
    }

    template <typename... Sizes>
    value_type operator()(size_t i, Sizes... /*sizes*/) const {
        return value_type(min_index(sub(i)));
    }

    template <typename E>
    bool alias(const E& rhs) const noexcept {
        return sub.alias(rhs);
    }

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(visitor));
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure sub value
        sub.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        sub.ensure_gpu_up_to_date();
    }

    friend std::ostream& operator<<(std::ostream& os, const argmin_transformer& transformer) {
        return os << "argmin(" << transformer.sub << ")";
    }
};

template <typename T>
struct sum_r_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    friend struct etl_traits<sum_r_transformer>;

    static constexpr bool gpu_computable = false;

private:
    sub_type sub; 

public:
    explicit sum_r_transformer(sub_type expr) : sub(expr) {}

    value_type operator[](size_t i) const {
        return sum(sub(i));
    }

    value_type read_flat(size_t i) const {
        return sum(sub(i));
    }

    template <typename... Sizes>
    value_type operator()(size_t i, Sizes... /*sizes*/) const {
        return sum(sub(i));
    }

    template <typename E>
    bool alias(const E& rhs) const noexcept {
        return sub.alias(rhs);
    }

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(visitor));
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure sub value
        sub.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        sub.ensure_gpu_up_to_date();
    }

    friend std::ostream& operator<<(std::ostream& os, const sum_r_transformer& transformer) {
        return os << "sum_r(" << transformer.sub << ")";
    }
};

template <typename T>
struct mean_r_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    friend struct etl_traits<mean_r_transformer>;

    static constexpr bool gpu_computable = false;

private:
    sub_type sub; 

public:
    explicit mean_r_transformer(sub_type expr) : sub(expr) {}

    value_type operator[](size_t i) const {
        return mean(sub(i));
    }

    value_type read_flat(size_t i) const {
        return mean(sub(i));
    }

    template <typename... Sizes>
    value_type operator()(size_t i, Sizes... /*sizes*/) const {
        return mean(sub(i));
    }

    template <typename E>
    bool alias(const E& rhs) const noexcept {
        return sub.alias(rhs);
    }

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(visitor));
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure sub value
        sub.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        sub.ensure_gpu_up_to_date();
    }

    friend std::ostream& operator<<(std::ostream& os, const mean_r_transformer& transformer) {
        return os << "mean_r(" << transformer.sub << ")";
    }
};

template <typename T>
struct sum_l_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    friend struct etl_traits<sum_l_transformer>;

    static constexpr bool gpu_computable = false;

private:
    sub_type sub; 

public:
    explicit sum_l_transformer(sub_type expr) : sub(expr) {}

    value_type operator[](size_t j) const {
        value_type m = 0.0;

        for (size_t i = 0; i < dim<0>(sub); ++i) {
            m += sub[j + i * (etl::size(sub) / dim<0>(sub))];
        }

        return m;
    }

    value_type read_flat(size_t j) const noexcept {
        value_type m = 0.0;

        for (size_t i = 0; i < dim<0>(sub); ++i) {
            m += sub.read_flat(j + i * (etl::size(sub) / dim<0>(sub)));
        }

        return m;
    }

    template <typename... Sizes>
    value_type operator()(size_t j, Sizes... sizes) const {
        value_type m = 0.0;

        for (size_t i = 0; i < dim<0>(sub); ++i) {
            m += sub(i, j, sizes...);
        }

        return m;
    }

    template <typename E>
    bool alias(const E& rhs) const noexcept {
        return sub.alias(rhs);
    }

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(visitor));
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure sub value
        sub.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        sub.ensure_gpu_up_to_date();
    }

    friend std::ostream& operator<<(std::ostream& os, const sum_l_transformer& transformer) {
        return os << "sum_l(" << transformer.sub << ")";
    }
};

template <typename T>
struct mean_l_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    friend struct etl_traits<mean_l_transformer>;

    static constexpr bool gpu_computable = false;

private:
    sub_type sub; 

public:
    explicit mean_l_transformer(sub_type expr) : sub(expr) {}

    value_type operator[](size_t j) const {
        value_type m = 0.0;

        for (size_t i = 0; i < dim<0>(sub); ++i) {
            m += sub[j + i * (etl::size(sub) / dim<0>(sub))];
        }

        return m / value_type(dim<0>(sub));
    }

    value_type read_flat(size_t j) const noexcept {
        value_type m = 0.0;

        for (size_t i = 0; i < dim<0>(sub); ++i) {
            m += sub.read_flat(j + i * (etl::size(sub) / dim<0>(sub)));
        }

        return m / value_type(dim<0>(sub));
    }

    template <typename... Sizes>
    value_type operator()(size_t j, Sizes... sizes) const {
        value_type m = 0.0;

        for (size_t i = 0; i < dim<0>(sub); ++i) {
            m += sub(i, j, sizes...);
        }

        return m / value_type(dim<0>(sub));
    }

    template <typename E>
    bool alias(const E& rhs) const noexcept {
        return sub.alias(rhs);
    }

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(visitor));
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure sub value
        sub.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        sub.ensure_gpu_up_to_date();
    }

    friend std::ostream& operator<<(std::ostream& os, const mean_l_transformer& transformer) {
        return os << "mean_l(" << transformer.sub << ")";
    }
};

template <typename T>
requires(
           cpp::specialization_of<etl::argmax_transformer, T> 
        || cpp::specialization_of<etl::argmin_transformer, T>
        || cpp::specialization_of<etl::sum_r_transformer, T>
        || cpp::specialization_of<etl::mean_r_transformer, T>)
struct etl_traits<T> {
    using expr_t     = T;                                                
    using sub_expr_t = std::decay_t<typename std::decay_t<T>::sub_type>; 
    using value_type = value_t<sub_expr_t>;                              

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

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

    static size_t size(const expr_t& v) {
        return etl::dim<0>(v.sub);
    }

    static size_t dim(const expr_t& v, [[maybe_unused]] size_t d) {
        return etl::dim<0>(v.sub);
    }

    static constexpr size_t size() {
        return etl_traits<sub_expr_t>::template dim<0>();
    }

    template <size_t D>
    static constexpr size_t dim() {
        return etl_traits<sub_expr_t>::template dim<0>();
    }

    static constexpr size_t dimensions() {
        return 1;
    }

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

template <typename T>
requires(cpp::specialization_of<etl::sum_l_transformer, T> || cpp::specialization_of<etl::mean_l_transformer, T>)
struct etl_traits<T> {
    using expr_t     = T;                                                
    using sub_expr_t = std::decay_t<typename std::decay_t<T>::sub_type>; 
    using value_type = value_t<sub_expr_t>;                              

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

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

    static size_t size(const expr_t& v) {
        return etl::size(v.sub) / etl::dim<0>(v.sub);
    }

    static size_t dim(const expr_t& v, size_t d) {
        return etl::dim(v.sub, d + 1);
    }

    static constexpr size_t size() {
        return etl_traits<sub_expr_t>::size() / etl_traits<sub_expr_t>::template dim<0>();
    }

    template <size_t D>
    static constexpr size_t dim() {
        return etl_traits<sub_expr_t>::template dim<D + 1>();
    }

    static constexpr size_t dimensions() {
        return etl_traits<sub_expr_t>::dimensions() - 1;
    }

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

} //end of namespace etl