batch_dim.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_BATCH_DIM_HPP_
#define GKO_PUBLIC_CORE_BASE_BATCH_DIM_HPP_
 
 
#include <iostream>
 
#include <ginkgo/core/base/dim.hpp>
#include <ginkgo/core/base/types.hpp>
 
 
namespace gko {
 
 
template <size_type Dimensionality = 2, typename DimensionType = size_type>
struct batch_dim {
    static constexpr size_type dimensionality = Dimensionality;
    using dimension_type = DimensionType;
 
    size_type get_num_batch_items() const { return num_batch_items_; }
 
    dim<dimensionality, dimension_type> get_common_size() const
    {
        return common_size_;
    }
 
    friend bool operator==(const batch_dim& x, const batch_dim& y)
    {
        return x.num_batch_items_ == y.num_batch_items_ &&
               x.common_size_ == y.common_size_;
    }
 
 
    friend bool operator!=(const batch_dim<Dimensionality, DimensionType>& x,
                           const batch_dim<Dimensionality, DimensionType>& y)
    {
        return !(x == y);
    }
 
 
    batch_dim()
        : common_size_(dim<dimensionality, dimension_type>{}),
          num_batch_items_(0)
    {}
 
    explicit batch_dim(const size_type num_batch_items,
                       const dim<dimensionality, dimension_type>& common_size)
        : common_size_(common_size), num_batch_items_(num_batch_items)
    {}
 
private:
    size_type num_batch_items_{};
    dim<dimensionality, dimension_type> common_size_{};
};
 
 
template <typename DimensionType>
inline batch_dim<2, DimensionType> transpose(
    const batch_dim<2, DimensionType>& input)
{
    return batch_dim<2, DimensionType>(input.get_num_batch_items(),
                                       transpose(input.get_common_size()));
}
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_BATCH_DIM_HPP_