mul.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef pRC_TENSOR_TRAIN_OPERATOR_FUNCTIONS_MUL_H
#define pRC_TENSOR_TRAIN_OPERATOR_FUNCTIONS_MUL_H
 
#include <prc/core/functors/view.hpp>
#include <prc/core/tensor/functions/reshape.hpp>
#include <prc/tensor_train/common/functions/loop.hpp>
 
namespace pRC
{
    template<class XA, class XB, class RA = RemoveReference<XA>,
        class RB = RemoveReference<XB>, If<TensorTrain::IsOperatorish<RA>> = 0,
        If<IsInvocable<View, XA>> = 0, If<TensorTrain::IsTensorish<RB>> = 0,
        If<IsInvocable<View, XB>> = 0,
        If<IsSame<typename RA::N, typename RB::N>> = 0>
    static inline constexpr auto operator*(XA &&a, XB &&b)
    {
        return loop(
            []<class XLA, class XLB>(XLA &&a, XLB &&b)
            {
                using CA = RemoveReference<XLA>;
                using CB = RemoveReference<XLB>;
 
                return reshape<CA::size(0) * CB::size(0), CA::size(1),
                    CA::size(3) * CB::size(2)>(permute<0, 3, 1, 2, 4>(
                    eval(contract<2, 1>(forward<XLA>(a), forward<XLB>(b)))));
            },
            forward<XA>(a), forward<XB>(b));
    }
 
    template<class XA, class XB, class RA = RemoveReference<XA>,
        class RB = RemoveReference<XB>, If<TensorTrain::IsTensorish<RA>> = 0,
        If<IsInvocable<View, XA>> = 0, If<TensorTrain::IsOperatorish<RB>> = 0,
        If<IsInvocable<View, XB>> = 0,
        If<IsSame<typename RA::N, typename RB::M>> = 0>
    static inline constexpr auto operator*(XA &&a, XB &&b)
    {
        return loop(
            []<class XLA, class XLB>(XLA &&a, XLB &&b)
            {
                using CA = RemoveReference<XLA>;
                using CB = RemoveReference<XLB>;
 
                return reshape<CA::size(0) * CB::size(0), CB::size(2),
                    CA::size(2) * CB::size(3)>(permute<0, 2, 3, 1, 4>(
                    eval(contract<1, 1>(forward<XLA>(a), forward<XLB>(b)))));
            },
            forward<XA>(a), forward<XB>(b));
    }
 
    template<class XA, class XB, class RA = RemoveReference<XA>,
        class RB = RemoveReference<XB>, If<TensorTrain::IsOperatorish<RA>> = 0,
        If<IsInvocable<View, XA>> = 0, If<TensorTrain::IsOperatorish<RB>> = 0,
        If<IsInvocable<View, XB>> = 0,
        If<IsSame<typename RA::N, typename RB::M>> = 0>
    static inline constexpr auto operator*(XA &&a, XB &&b)
    {
        return loop(
            []<class XLA, class XLB>(XLA &&a, XLB &&b)
            {
                using CA = RemoveReference<XLA>;
                using CB = RemoveReference<XLB>;
 
                return reshape<CA::size(0) * CB::size(0), CA::size(1),
                    CB::size(2), CA::size(3) * CB::size(3)>(
                    permute<0, 3, 1, 4, 2, 5>(eval(
                        contract<2, 1>(forward<XLA>(a), forward<XLB>(b)))));
            },
            forward<XA>(a), forward<XB>(b));
    }
 
    template<class XA, class XB, class RA = RemoveReference<XA>,
        class RB = RemoveReference<XB>,
        If<Any<TensorTrain::IsOperatorish<RA>, TensorTrain::IsTensorish<RA>>> =
            0,
        If<Any<TensorTrain::IsOperatorish<RB>, TensorTrain::IsTensorish<RB>>> =
            0,
        If<Not<All<IsInvocable<View, XA>, IsInvocable<View, XB>>>> = 0,
        If<IsInvocable<Mul, XA &, XB &>> = 0>
    static inline constexpr auto operator*(XA &&a, XB &&b)
    {
        return eval(a * b);
    }
}
#endif // pRC_TENSOR_TRAIN_OPERATOR_FUNCTIONS_MUL_H