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/tensor_train/operator/declarations.hpp>
#include <prc/tensor_train/tensor/declarations.hpp>
 
namespace pRC
{
    template<class XA, class XB,
        TensorTrain::IsOperatorish RA = RemoveReference<XA>,
        TensorTrain::IsTensorish RB = RemoveReference<XB>>
        requires IsSame<typename RA::N, typename RB::N>
    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,
        TensorTrain::IsTensorish RA = RemoveReference<XA>,
        TensorTrain::IsOperatorish RB = RemoveReference<XB>>
        requires IsSame<typename RA::N, typename RB::M>
    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,
        TensorTrain::IsOperatorish RA = RemoveReference<XA>,
        TensorTrain::IsOperatorish RB = RemoveReference<XB>>
        requires IsSame<typename RA::N, typename RB::M>
    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));
    }
}
#endif // pRC_TENSOR_TRAIN_OPERATOR_FUNCTIONS_MUL_H