shape_utils.h

// Copyright 2024, UChicago Argonne, LLC
// All Rights Reserved
// Software Name: NEML2 -- the New Engineering material Model Library, version 2
// By: Argonne National Laboratory
// OPEN SOURCE LICENSE (MIT)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
 
#pragma once
 
#include "neml2/misc/types.h"
#include "neml2/misc/errors.h"
 
namespace neml2::utils
{
template <class... T>
bool broadcastable(const T &... tensors);
 
template <class... T>
bool batch_broadcastable(const T &... tensors);
 
template <class... T>
bool base_broadcastable(const T &... tensors);
 
template <class... T>
Size broadcast_batch_dim(const T &...);
 
template <class... T>
bool sizes_same(T &&... shapes);
 
template <class... T>
bool sizes_broadcastable(const T &... shapes);
 
template <class... T>
TensorShape broadcast_sizes(const T &... shapes);
 
Size storage_size(TensorShapeRef shape);
 
template <typename... S>
TensorShape add_shapes(const S &...);
 
TensorShape pad_prepend(TensorShapeRef s, Size dim, Size pad = 1);
 
namespace details
{
template <typename... S>
TensorShape add_shapes_impl(TensorShape &, TensorShapeRef, const S &...);
} // namespace details
} // namespace neml2::utils
 
// Implementation
 
namespace neml2::utils
{
template <class... T>
bool
broadcastable(const T &... tensors)
{
  if (!sizes_same(tensors.base_sizes()...))
    return false;
  return batch_broadcastable(tensors...);
}
 
template <class... T>
bool
batch_broadcastable(const T &... tensors)
{
  return sizes_broadcastable(tensors.batch_sizes().concrete()...);
}
 
template <class... T>
bool
base_broadcastable(const T &... tensors)
{
  return sizes_broadcastable(tensors.base_sizes()...);
}
 
template <class... T>
Size
broadcast_batch_dim(const T &... tensor)
{
  return std::max({tensor.batch_dim()...});
}
 
template <class... T>
bool
sizes_same(T &&... shapes)
{
  auto all_shapes = std::vector<TensorShapeRef>{std::forward<T>(shapes)...};
  for (size_t i = 0; i < all_shapes.size() - 1; i++)
    if (all_shapes[i] != all_shapes[i + 1])
      return false;
  return true;
}
 
template <class... T>
bool
sizes_broadcastable(const T &... shapes)
{
  auto dim = std::max({shapes.size()...});
  auto all_shapes_padded = std::vector<TensorShape>{pad_prepend(shapes, dim)...};
 
  for (size_t i = 0; i < dim; i++)
  {
    Size max_sz = 1;
    for (const auto & s : all_shapes_padded)
    {
      if (max_sz == 1)
      {
#ifndef NDEBUG
        if (s[i] <= 0)
          throw NEMLException("Found a size equal or less than 0: " + std::to_string(s[i]));
#endif
        if (s[i] > 1)
          max_sz = s[i];
      }
      else if (s[i] != 1 && s[i] != max_sz)
        return false;
    }
  }
 
  return true;
}
 
template <class... T>
TensorShape
broadcast_sizes(const T &... shapes)
{
#ifndef NDEBUG
  if (!sizes_broadcastable(shapes...))
    throw NEMLException("Shapes not broadcastable");
#endif
 
  auto dim = std::max({shapes.size()...});
  auto all_shapes_padded = std::vector<TensorShape>{pad_prepend(shapes, dim)...};
  auto bshape = TensorShape(dim, 1);
 
  for (size_t i = 0; i < dim; i++)
    for (const auto & s : all_shapes_padded)
      if (s[i] > bshape[i])
        bshape[i] = s[i];
 
  return bshape;
}
 
template <typename... S>
TensorShape
add_shapes(const S &... shape)
{
  TensorShape net;
  return details::add_shapes_impl(net, shape...);
}
 
namespace details
{
template <typename... S>
TensorShape
add_shapes_impl(TensorShape & net, TensorShapeRef s, const S &... rest)
{
  net.insert(net.end(), s.begin(), s.end());
 
  if constexpr (sizeof...(rest) == 0)
    return std::move(net);
  else
    return add_shapes_impl(net, rest...);
}
} // namespace details
} // namespace neml2::utils