add.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef pRC_TENSOR_TRAIN_TENSOR_FUNCTIONS_ADD_H
#define pRC_TENSOR_TRAIN_TENSOR_FUNCTIONS_ADD_H
 
#include <prc/core/functors/direct_sum.hpp>
#include <prc/tensor_train/tensor/declarations.hpp>
 
namespace pRC
{
    template<class XA, class XB,
        TensorTrain::IsTensorish RA = RemoveReference<XA>,
        TensorTrain::IsTensorish RB = RemoveReference<XB>>
        requires IsSame<typename RA::Sizes, typename RB::Sizes>
    static inline constexpr auto operator+(XA &&a, XB &&b)
    {
        return enumerate(
            []<Index C>(auto &&...args) -> decltype(auto)
            {
                if constexpr(C == 0)
                {
                    return exclude<DirectSum, 0, 1>(
                        forward<decltype(args)>(args)...);
                }
                else if constexpr(C == RA::Dimension - 1 &&
                    C == RB::Dimension - 1)
                {
                    return permute<2, 0, 1>(exclude<DirectSum, 1, 2>(
                        forward<decltype(args)>(args)...));
                }
                else
                {
                    return permute<1, 0, 2>(exclude<DirectSum, 1>(
                        forward<decltype(args)>(args)...));
                }
            },
            forward<XA>(a), forward<XB>(b));
    }
}
#endif // pRC_TENSOR_TRAIN_TENSOR_FUNCTIONS_ADD_H