matrix_product_shape.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/common/trace.hpp>
#include <uni20/core/types.hpp>
#include <uni20/mdspan/mdspan.hpp>
#include <uni20/tensor/concepts.hpp>
 
namespace uni20::linalg::detail
{
 
using matrix_product_extents = stdex::dextents<uni20::index_type, 2>;
 
template <class Lhs, class Rhs>
[[nodiscard]] matrix_product_extents matrix_product_shape_impl(Lhs const& lhs, Rhs const& rhs)
{
  ERROR_IF(lhs.extent(1) != rhs.extent(0), "matrix product inner extents do not agree", lhs.extent(1), rhs.extent(0));
  return matrix_product_extents{static_cast<uni20::index_type>(lhs.extent(0)),
                                static_cast<uni20::index_type>(rhs.extent(1))};
}
 
template <uni20::RankedTensorView<2> LhsTensor, uni20::RankedTensorView<2> RhsTensor>
[[nodiscard]] matrix_product_extents matrix_product_shape(LhsTensor const& lhs, RhsTensor const& rhs)
{
  return matrix_product_shape_impl(lhs, rhs);
}
 
template <uni20::RankedMdspecLike<2> LhsMdspan, uni20::RankedMdspecLike<2> RhsMdspan>
[[nodiscard]] matrix_product_extents matrix_product_shape(LhsMdspan const& lhs, RhsMdspan const& rhs)
{
  return matrix_product_shape_impl(lhs, rhs);
}
 
} // namespace uni20::linalg::detail