blocktensor.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <array>
#include <exception>
#include <initializer_list>
#include <memory>
#include <type_traits>
 
#include <core.hpp>
#include <utils/fqn.hpp>
 
namespace iganet {
namespace utils {
 
template <typename T> struct is_shared_ptr : std::false_type {};
 
template <typename T>
struct is_shared_ptr<std::shared_ptr<T>> : std::true_type {};
 
template <typename T> inline auto make_shared(T &&arg) {
  if constexpr (is_shared_ptr<typename std::decay<T>::type>::value)
    return std::forward<typename std::decay<T>::type>(arg);
  else
    return std::make_shared<typename std::decay<T>::type>(std::forward<T>(arg));
}
 
template <typename T, std::size_t... Dims> class BlockTensor;
 
template <typename T, std::size_t... Dims>
class BlockTensorCore : protected iganet::utils::FullQualifiedName {
 
protected:
  std::array<std::shared_ptr<T>, (Dims * ...)> data_;
 
public:
  BlockTensorCore() = default;
 
  template <typename... Ts, std::size_t... dims>
  BlockTensorCore(BlockTensorCore<Ts, dims...> &&...other) {
    auto it = data_.begin();
    (std::transform(other.data().begin(), other.data().end(), it,
                    [&it](auto &&d) {
                      it++;
                      return std::move(d);
                    }),
     ...);
  }
 
  template <typename... Ts, std::size_t... dims>
  BlockTensorCore(BlockTensor<Ts, dims...> &&...other) {
    auto it = data_.begin();
    (std::transform(other.data().begin(), other.data().end(), it,
                    [&it](auto &&d) {
                      it++;
                      return std::move(d);
                    }),
     ...);
  }
 
  template <typename... Ts>
  explicit BlockTensorCore(Ts &&...data)
      : data_({make_shared<Ts>(std::forward<Ts>(data))...}) {}
 
  inline static constexpr auto dims() {
    return std::array<std::size_t, sizeof...(Dims)>({Dims...});
  }
 
  template <std::size_t i> inline static constexpr std::size_t dim() {
    if constexpr (i < sizeof...(Dims))
      return std::get<i>(std::forward_as_tuple(Dims...));
    else
      return 0;
  }
 
  inline static constexpr std::size_t size() { return sizeof...(Dims); }
 
  inline static constexpr std::size_t entries() { return (Dims * ...); }
 
  inline const std::array<std::shared_ptr<T>, (Dims * ...)> &data() const {
    return data_;
  }
 
  inline std::array<std::shared_ptr<T>, (Dims * ...)> &data() { return data_; }
 
  inline const std::shared_ptr<T> &operator[](std::size_t idx) const {
    assert(0 <= idx && idx < (Dims * ...));
    return data_[idx];
  }
 
  inline std::shared_ptr<T> &operator[](std::size_t idx) {
    assert(0 <= idx && idx < (Dims * ...));
    return data_[idx];
  }
 
  inline const T &operator()(std::size_t idx) const {
    assert(0 <= idx && idx < (Dims * ...));
    return *data_[idx];
  }
 
  inline T &operator()(std::size_t idx) {
    assert(0 <= idx && idx < (Dims * ...));
    return *data_[idx];
  }
 
  template <typename Data> inline T &set(std::size_t idx, Data &&data) {
    assert(0 <= idx && idx < (Dims * ...));
    data_[idx] = make_shared<Data>(std::forward<Data>(data));
    return *data_[idx];
  }
 
  inline virtual void
  pretty_print(std::ostream &os = Log(log::info)) const noexcept = 0;
};
 
template <typename T, std::size_t... Dims>
inline std::ostream &operator<<(std::ostream &os,
                                const BlockTensorCore<T, Dims...> &obj) {
  obj.pretty_print(os);
  return os;
}
 
template <typename T, std::size_t Rows>
class BlockTensor<T, Rows> : public BlockTensorCore<T, Rows> {
private:
  using Base = BlockTensorCore<T, Rows>;
 
public:
  using BlockTensorCore<T, Rows>::BlockTensorCore;
 
  inline static constexpr std::size_t rows() { return Rows; }
 
  inline virtual void
  pretty_print(std::ostream &os = Log(log::info)) const noexcept override {
    os << Base::name() << "\n";
    for (std::size_t row = 0; row < Rows; ++row)
      os << "[" << row << "] = \n" << *Base::data_[row] << "\n";
  }
};
 
template <typename T, std::size_t Rows, std::size_t Cols>
class BlockTensor<T, Rows, Cols> : public BlockTensorCore<T, Rows, Cols> {
private:
  using Base = BlockTensorCore<T, Rows, Cols>;
 
public:
  using BlockTensorCore<T, Rows, Cols>::BlockTensorCore;
 
  inline static constexpr std::size_t rows() { return Rows; }
 
  inline static constexpr std::size_t cols() { return Cols; }
 
  using Base::operator();
 
  inline const T &operator()(std::size_t row, std::size_t col) const {
    assert(0 <= row && row < Rows && 0 <= col && col < Cols);
    return *Base::data_[Cols * row + col];
  }
 
  inline T &operator()(std::size_t row, std::size_t col) {
    assert(0 <= row && row < Rows && 0 <= col && col < Cols);
    return *Base::data_[Cols * row + col];
  }
 
  using Base::set;
 
  template <typename D>
  inline T &set(std::size_t row, std::size_t col, D &&data) {
    assert(0 <= row && row < Rows && 0 <= col && col < Cols);
    Base::data_[Cols * row + col] = make_shared<D>(std::forward<D>(data));
    return *Base::data_[Cols * row + col];
  }
 
  inline auto tr() const {
    BlockTensor<T, Cols, Rows> result;
    for (std::size_t row = 0; row < Rows; ++row)
      for (std::size_t col = 0; col < Cols; ++col)
        result[Rows * col + row] = Base::data_[Cols * row + col];
    return result;
  }
 
 
  inline auto det() const {
    if constexpr (Rows == 1 && Cols == 1) {
      auto result = *Base::data_[0];
      return result;
    } else if constexpr (Rows == 2 && Cols == 2) {
      auto result = torch::mul(*Base::data_[0], *Base::data_[3]) -
                    torch::mul(*Base::data_[1], *Base::data_[2]);
      return result;
    } else if constexpr (Rows == 3 && Cols == 3) {
      auto result =
          torch::mul(*Base::data_[0],
                     torch::mul(*Base::data_[4], *Base::data_[8]) -
                         torch::mul(*Base::data_[5], *Base::data_[7])) -
          torch::mul(*Base::data_[1],
                     torch::mul(*Base::data_[3], *Base::data_[8]) -
                         torch::mul(*Base::data_[5], *Base::data_[6])) +
          torch::mul(*Base::data_[2],
                     torch::mul(*Base::data_[3], *Base::data_[7]) -
                         torch::mul(*Base::data_[4], *Base::data_[6]));
      return result;
    } else if constexpr (Rows == 4 && Cols == 4) {
      auto a11 = torch::mul(*Base::data_[5],
                            (torch::mul(*Base::data_[10], *Base::data_[15]) -
                             torch::mul(*Base::data_[11], *Base::data_[14]))) -
                 torch::mul(*Base::data_[9],
                            (torch::mul(*Base::data_[6], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[14]))) -
                 torch::mul(*Base::data_[13],
                            (torch::mul(*Base::data_[7], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[11])));
 
      auto a21 = torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[11], *Base::data_[14]) -
                             torch::mul(*Base::data_[10], *Base::data_[15]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[7], *Base::data_[14]) -
                             torch::mul(*Base::data_[6], *Base::data_[15]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[6], *Base::data_[11]) -
                             torch::mul(*Base::data_[7], *Base::data_[10])));
 
      auto a31 = torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[9], *Base::data_[15]) -
                             torch::mul(*Base::data_[11], *Base::data_[13]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[5], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[13]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[7], *Base::data_[9]) -
                             torch::mul(*Base::data_[5], *Base::data_[11])));
 
      auto a41 = torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[10], *Base::data_[13]) -
                             torch::mul(*Base::data_[9], *Base::data_[14]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[6], *Base::data_[13]) -
                             torch::mul(*Base::data_[5], *Base::data_[14]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[5], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[9])));
 
      auto result =
          torch::mul(*Base::data_[0], a11) + torch::mul(*Base::data_[1], a21) +
          torch::mul(*Base::data_[2], a31) + torch::mul(*Base::data_[3], a41);
 
      return result;
    } else {
      throw std::runtime_error("Unsupported block tensor dimension");
      return *this;
    }
  }
 
  inline auto inv() const {
 
    auto det_ = this->det();
 
    if constexpr (Rows == 1 && Cols == 1) {
      BlockTensor<T, Rows, Cols> result;
      result[0] = std::make_shared<T>(torch::reciprocal(*Base::data_[0]));
      return result;
    } else if constexpr (Rows == 2 && Cols == 2) {
 
      BlockTensor<T, Rows, Cols> result;
      result[0] = std::make_shared<T>(torch::div(*Base::data_[3], det_));
      result[1] = std::make_shared<T>(torch::div(*Base::data_[2], -det_));
      result[2] = std::make_shared<T>(torch::div(*Base::data_[1], -det_));
      result[3] = std::make_shared<T>(torch::div(*Base::data_[0], det_));
      return result;
    } else if constexpr (Rows == 3 && Cols == 3) {
 
      auto a11 = torch::mul(*Base::data_[4], *Base::data_[8]) -
                 torch::mul(*Base::data_[5], *Base::data_[7]);
      auto a12 = torch::mul(*Base::data_[2], *Base::data_[7]) -
                 torch::mul(*Base::data_[1], *Base::data_[8]);
      auto a13 = torch::mul(*Base::data_[1], *Base::data_[5]) -
                 torch::mul(*Base::data_[2], *Base::data_[4]);
      auto a21 = torch::mul(*Base::data_[5], *Base::data_[6]) -
                 torch::mul(*Base::data_[3], *Base::data_[8]);
      auto a22 = torch::mul(*Base::data_[0], *Base::data_[8]) -
                 torch::mul(*Base::data_[2], *Base::data_[6]);
      auto a23 = torch::mul(*Base::data_[2], *Base::data_[3]) -
                 torch::mul(*Base::data_[0], *Base::data_[5]);
      auto a31 = torch::mul(*Base::data_[3], *Base::data_[7]) -
                 torch::mul(*Base::data_[4], *Base::data_[6]);
      auto a32 = torch::mul(*Base::data_[1], *Base::data_[6]) -
                 torch::mul(*Base::data_[0], *Base::data_[7]);
      auto a33 = torch::mul(*Base::data_[0], *Base::data_[4]) -
                 torch::mul(*Base::data_[1], *Base::data_[3]);
 
      BlockTensor<T, Rows, Cols> result;
      result[0] = std::make_shared<T>(torch::div(a11, det_));
      result[1] = std::make_shared<T>(torch::div(a12, det_));
      result[2] = std::make_shared<T>(torch::div(a13, det_));
      result[3] = std::make_shared<T>(torch::div(a21, det_));
      result[4] = std::make_shared<T>(torch::div(a22, det_));
      result[5] = std::make_shared<T>(torch::div(a23, det_));
      result[6] = std::make_shared<T>(torch::div(a31, det_));
      result[7] = std::make_shared<T>(torch::div(a32, det_));
      result[8] = std::make_shared<T>(torch::div(a33, det_));
      return result;
    } else if constexpr (Rows == 4 && Cols == 4) {
      auto a11 = torch::mul(*Base::data_[5],
                            (torch::mul(*Base::data_[10], *Base::data_[15]) -
                             torch::mul(*Base::data_[11], *Base::data_[14]))) -
                 torch::mul(*Base::data_[9],
                            (torch::mul(*Base::data_[6], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[14]))) -
                 torch::mul(*Base::data_[13],
                            (torch::mul(*Base::data_[7], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[11])));
 
      auto a12 = torch::mul(*Base::data_[1],
                            (torch::mul(*Base::data_[11], *Base::data_[14]) -
                             torch::mul(*Base::data_[10], *Base::data_[15]))) -
                 torch::mul(*Base::data_[9],
                            (torch::mul(*Base::data_[3], *Base::data_[14]) -
                             torch::mul(*Base::data_[2], *Base::data_[15]))) -
                 torch::mul(*Base::data_[13],
                            (torch::mul(*Base::data_[2], *Base::data_[11]) -
                             torch::mul(*Base::data_[3], *Base::data_[10])));
 
      auto a13 = torch::mul(*Base::data_[1],
                            (torch::mul(*Base::data_[6], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[14]))) -
                 torch::mul(*Base::data_[5],
                            (torch::mul(*Base::data_[2], *Base::data_[15]) -
                             torch::mul(*Base::data_[3], *Base::data_[14]))) -
                 torch::mul(*Base::data_[13],
                            (torch::mul(*Base::data_[3], *Base::data_[6]) -
                             torch::mul(*Base::data_[2], *Base::data_[7])));
 
      auto a14 = torch::mul(*Base::data_[1],
                            (torch::mul(*Base::data_[7], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[11]))) -
                 torch::mul(*Base::data_[5],
                            (torch::mul(*Base::data_[3], *Base::data_[10]) -
                             torch::mul(*Base::data_[2], *Base::data_[11]))) -
                 torch::mul(*Base::data_[9],
                            (torch::mul(*Base::data_[2], *Base::data_[7]) -
                             torch::mul(*Base::data_[3], *Base::data_[6])));
 
      auto a21 = torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[11], *Base::data_[14]) -
                             torch::mul(*Base::data_[10], *Base::data_[15]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[7], *Base::data_[14]) -
                             torch::mul(*Base::data_[6], *Base::data_[15]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[6], *Base::data_[11]) -
                             torch::mul(*Base::data_[7], *Base::data_[10])));
 
      auto a22 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[10], *Base::data_[15]) -
                             torch::mul(*Base::data_[11], *Base::data_[14]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[2], *Base::data_[15]) -
                             torch::mul(*Base::data_[3], *Base::data_[14]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[3], *Base::data_[10]) -
                             torch::mul(*Base::data_[2], *Base::data_[11])));
 
      auto a23 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[7], *Base::data_[14]) -
                             torch::mul(*Base::data_[6], *Base::data_[15]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[3], *Base::data_[14]) -
                             torch::mul(*Base::data_[2], *Base::data_[15]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[2], *Base::data_[7]) -
                             torch::mul(*Base::data_[3], *Base::data_[6])));
 
      auto a24 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[6], *Base::data_[11]) -
                             torch::mul(*Base::data_[7], *Base::data_[10]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[2], *Base::data_[11]) -
                             torch::mul(*Base::data_[3], *Base::data_[10]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[3], *Base::data_[6]) -
                             torch::mul(*Base::data_[2], *Base::data_[7])));
 
      auto a31 = torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[9], *Base::data_[15]) -
                             torch::mul(*Base::data_[11], *Base::data_[13]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[5], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[13]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[7], *Base::data_[9]) -
                             torch::mul(*Base::data_[5], *Base::data_[11])));
 
      auto a32 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[11], *Base::data_[13]) -
                             torch::mul(*Base::data_[9], *Base::data_[15]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[3], *Base::data_[13]) -
                             torch::mul(*Base::data_[1], *Base::data_[15]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[1], *Base::data_[11]) -
                             torch::mul(*Base::data_[3], *Base::data_[9])));
 
      auto a33 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[5], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[13]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[1], *Base::data_[15]) -
                             torch::mul(*Base::data_[3], *Base::data_[13]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[3], *Base::data_[5]) -
                             torch::mul(*Base::data_[1], *Base::data_[7])));
 
      auto a34 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[7], *Base::data_[9]) -
                             torch::mul(*Base::data_[5], *Base::data_[11]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[3], *Base::data_[9]) -
                             torch::mul(*Base::data_[1], *Base::data_[11]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[1], *Base::data_[7]) -
                             torch::mul(*Base::data_[3], *Base::data_[5])));
 
      auto a41 = torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[10], *Base::data_[13]) -
                             torch::mul(*Base::data_[9], *Base::data_[14]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[6], *Base::data_[13]) -
                             torch::mul(*Base::data_[5], *Base::data_[14]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[5], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[9])));
 
      auto a42 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[9], *Base::data_[14]) -
                             torch::mul(*Base::data_[10], *Base::data_[13]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[1], *Base::data_[14]) -
                             torch::mul(*Base::data_[2], *Base::data_[13]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[2], *Base::data_[9]) -
                             torch::mul(*Base::data_[1], *Base::data_[10])));
 
      auto a43 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[6], *Base::data_[13]) -
                             torch::mul(*Base::data_[5], *Base::data_[14]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[2], *Base::data_[13]) -
                             torch::mul(*Base::data_[1], *Base::data_[14]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[1], *Base::data_[6]) -
                             torch::mul(*Base::data_[2], *Base::data_[5])));
 
      auto a44 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[5], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[9]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[1], *Base::data_[10]) -
                             torch::mul(*Base::data_[2], *Base::data_[9]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[2], *Base::data_[5]) -
                             torch::mul(*Base::data_[1], *Base::data_[6])));
      BlockTensor<T, Rows, Cols> result;
      result[0] = std::make_shared<T>(torch::div(a11, det_));
      result[1] = std::make_shared<T>(torch::div(a12, det_));
      result[2] = std::make_shared<T>(torch::div(a13, det_));
      result[3] = std::make_shared<T>(torch::div(a14, det_));
      result[4] = std::make_shared<T>(torch::div(a21, det_));
      result[5] = std::make_shared<T>(torch::div(a22, det_));
      result[6] = std::make_shared<T>(torch::div(a23, det_));
      result[7] = std::make_shared<T>(torch::div(a24, det_));
      result[8] = std::make_shared<T>(torch::div(a31, det_));
      result[9] = std::make_shared<T>(torch::div(a32, det_));
      result[10] = std::make_shared<T>(torch::div(a33, det_));
      result[11] = std::make_shared<T>(torch::div(a34, det_));
      result[12] = std::make_shared<T>(torch::div(a41, det_));
      result[13] = std::make_shared<T>(torch::div(a42, det_));
      result[14] = std::make_shared<T>(torch::div(a43, det_));
      result[15] = std::make_shared<T>(torch::div(a44, det_));
      return result;
    } else {
      throw std::runtime_error("Unsupported block tensor dimension");
      return *this;
    }
  }
 
  inline auto ginv() const {
    if constexpr (Rows == Cols)
      return this->inv();
    else
      // Compute the generalized inverse, i.e. (A^T A)^{-1} A^T
      return (this->tr() * (*this)).inv() * this->tr();
  }
 
  inline auto invtr() const {
 
    auto det_ = this->det();
 
    if constexpr (Rows == 1 && Cols == 1) {
      BlockTensor<T, Cols, Rows> result;
      result[0] = std::make_shared<T>(torch::reciprocal(*Base::data_[0]));
      return result;
    } else if constexpr (Rows == 2 && Cols == 2) {
 
      BlockTensor<T, Cols, Rows> result;
      result[0] = std::make_shared<T>(torch::div(*Base::data_[3], det_));
      result[1] = std::make_shared<T>(torch::div(*Base::data_[1], -det_));
      result[2] = std::make_shared<T>(torch::div(*Base::data_[2], -det_));
      result[3] = std::make_shared<T>(torch::div(*Base::data_[0], det_));
      return result;
    } else if constexpr (Rows == 3 && Cols == 3) {
 
      auto a11 = torch::mul(*Base::data_[4], *Base::data_[8]) -
                 torch::mul(*Base::data_[5], *Base::data_[7]);
      auto a12 = torch::mul(*Base::data_[2], *Base::data_[7]) -
                 torch::mul(*Base::data_[1], *Base::data_[8]);
      auto a13 = torch::mul(*Base::data_[1], *Base::data_[5]) -
                 torch::mul(*Base::data_[2], *Base::data_[4]);
      auto a21 = torch::mul(*Base::data_[5], *Base::data_[6]) -
                 torch::mul(*Base::data_[3], *Base::data_[8]);
      auto a22 = torch::mul(*Base::data_[0], *Base::data_[8]) -
                 torch::mul(*Base::data_[2], *Base::data_[6]);
      auto a23 = torch::mul(*Base::data_[2], *Base::data_[3]) -
                 torch::mul(*Base::data_[0], *Base::data_[5]);
      auto a31 = torch::mul(*Base::data_[3], *Base::data_[7]) -
                 torch::mul(*Base::data_[4], *Base::data_[6]);
      auto a32 = torch::mul(*Base::data_[1], *Base::data_[6]) -
                 torch::mul(*Base::data_[0], *Base::data_[7]);
      auto a33 = torch::mul(*Base::data_[0], *Base::data_[4]) -
                 torch::mul(*Base::data_[1], *Base::data_[3]);
 
      BlockTensor<T, Cols, Rows> result;
      result[0] = std::make_shared<T>(torch::div(a11, det_));
      result[1] = std::make_shared<T>(torch::div(a21, det_));
      result[2] = std::make_shared<T>(torch::div(a31, det_));
      result[3] = std::make_shared<T>(torch::div(a12, det_));
      result[4] = std::make_shared<T>(torch::div(a22, det_));
      result[5] = std::make_shared<T>(torch::div(a32, det_));
      result[6] = std::make_shared<T>(torch::div(a13, det_));
      result[7] = std::make_shared<T>(torch::div(a23, det_));
      result[8] = std::make_shared<T>(torch::div(a33, det_));
      return result;
    } else if constexpr (Rows == 4 && Cols == 4) {
 
      auto a11 = torch::mul(*Base::data_[5],
                            (torch::mul(*Base::data_[10], *Base::data_[15]) -
                             torch::mul(*Base::data_[11], *Base::data_[14]))) -
                 torch::mul(*Base::data_[9],
                            (torch::mul(*Base::data_[6], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[14]))) -
                 torch::mul(*Base::data_[13],
                            (torch::mul(*Base::data_[7], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[11])));
 
      auto a12 = torch::mul(*Base::data_[1],
                            (torch::mul(*Base::data_[11], *Base::data_[14]) -
                             torch::mul(*Base::data_[10], *Base::data_[15]))) -
                 torch::mul(*Base::data_[9],
                            (torch::mul(*Base::data_[3], *Base::data_[14]) -
                             torch::mul(*Base::data_[2], *Base::data_[15]))) -
                 torch::mul(*Base::data_[13],
                            (torch::mul(*Base::data_[2], *Base::data_[11]) -
                             torch::mul(*Base::data_[3], *Base::data_[10])));
 
      auto a13 = torch::mul(*Base::data_[1],
                            (torch::mul(*Base::data_[6], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[14]))) -
                 torch::mul(*Base::data_[5],
                            (torch::mul(*Base::data_[2], *Base::data_[15]) -
                             torch::mul(*Base::data_[3], *Base::data_[14]))) -
                 torch::mul(*Base::data_[13],
                            (torch::mul(*Base::data_[3], *Base::data_[6]) -
                             torch::mul(*Base::data_[2], *Base::data_[7])));
 
      auto a14 = torch::mul(*Base::data_[1],
                            (torch::mul(*Base::data_[7], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[11]))) -
                 torch::mul(*Base::data_[5],
                            (torch::mul(*Base::data_[3], *Base::data_[10]) -
                             torch::mul(*Base::data_[2], *Base::data_[11]))) -
                 torch::mul(*Base::data_[9],
                            (torch::mul(*Base::data_[2], *Base::data_[7]) -
                             torch::mul(*Base::data_[3], *Base::data_[6])));
 
      auto a21 = torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[11], *Base::data_[14]) -
                             torch::mul(*Base::data_[10], *Base::data_[15]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[7], *Base::data_[14]) -
                             torch::mul(*Base::data_[6], *Base::data_[15]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[6], *Base::data_[11]) -
                             torch::mul(*Base::data_[7], *Base::data_[10])));
 
      auto a22 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[10], *Base::data_[15]) -
                             torch::mul(*Base::data_[11], *Base::data_[14]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[2], *Base::data_[15]) -
                             torch::mul(*Base::data_[3], *Base::data_[14]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[3], *Base::data_[10]) -
                             torch::mul(*Base::data_[2], *Base::data_[11])));
 
      auto a23 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[7], *Base::data_[14]) -
                             torch::mul(*Base::data_[6], *Base::data_[15]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[3], *Base::data_[14]) -
                             torch::mul(*Base::data_[2], *Base::data_[15]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[2], *Base::data_[7]) -
                             torch::mul(*Base::data_[3], *Base::data_[6])));
 
      auto a24 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[6], *Base::data_[11]) -
                             torch::mul(*Base::data_[7], *Base::data_[10]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[2], *Base::data_[11]) -
                             torch::mul(*Base::data_[3], *Base::data_[10]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[3], *Base::data_[6]) -
                             torch::mul(*Base::data_[2], *Base::data_[7])));
 
      auto a31 = torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[9], *Base::data_[15]) -
                             torch::mul(*Base::data_[11], *Base::data_[13]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[5], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[13]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[7], *Base::data_[9]) -
                             torch::mul(*Base::data_[5], *Base::data_[11])));
 
      auto a32 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[11], *Base::data_[13]) -
                             torch::mul(*Base::data_[9], *Base::data_[15]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[3], *Base::data_[13]) -
                             torch::mul(*Base::data_[1], *Base::data_[15]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[1], *Base::data_[11]) -
                             torch::mul(*Base::data_[3], *Base::data_[9])));
 
      auto a33 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[5], *Base::data_[15]) -
                             torch::mul(*Base::data_[7], *Base::data_[13]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[1], *Base::data_[15]) -
                             torch::mul(*Base::data_[3], *Base::data_[13]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[3], *Base::data_[5]) -
                             torch::mul(*Base::data_[1], *Base::data_[7])));
 
      auto a34 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[7], *Base::data_[9]) -
                             torch::mul(*Base::data_[5], *Base::data_[11]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[3], *Base::data_[9]) -
                             torch::mul(*Base::data_[1], *Base::data_[11]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[1], *Base::data_[7]) -
                             torch::mul(*Base::data_[3], *Base::data_[5])));
 
      auto a41 = torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[10], *Base::data_[13]) -
                             torch::mul(*Base::data_[9], *Base::data_[14]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[6], *Base::data_[13]) -
                             torch::mul(*Base::data_[5], *Base::data_[14]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[5], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[9])));
 
      auto a42 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[9], *Base::data_[14]) -
                             torch::mul(*Base::data_[10], *Base::data_[13]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[1], *Base::data_[14]) -
                             torch::mul(*Base::data_[2], *Base::data_[13]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[2], *Base::data_[9]) -
                             torch::mul(*Base::data_[1], *Base::data_[10])));
 
      auto a43 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[6], *Base::data_[13]) -
                             torch::mul(*Base::data_[5], *Base::data_[14]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[2], *Base::data_[13]) -
                             torch::mul(*Base::data_[1], *Base::data_[14]))) -
                 torch::mul(*Base::data_[12],
                            (torch::mul(*Base::data_[1], *Base::data_[6]) -
                             torch::mul(*Base::data_[2], *Base::data_[5])));
 
      auto a44 = torch::mul(*Base::data_[0],
                            (torch::mul(*Base::data_[5], *Base::data_[10]) -
                             torch::mul(*Base::data_[6], *Base::data_[9]))) -
                 torch::mul(*Base::data_[4],
                            (torch::mul(*Base::data_[1], *Base::data_[10]) -
                             torch::mul(*Base::data_[2], *Base::data_[9]))) -
                 torch::mul(*Base::data_[8],
                            (torch::mul(*Base::data_[2], *Base::data_[5]) -
                             torch::mul(*Base::data_[1], *Base::data_[6])));
 
      BlockTensor<T, Rows, Cols> result;
      result[0] = std::make_shared<T>(torch::div(a11, det_));
      result[1] = std::make_shared<T>(torch::div(a21, det_));
      result[2] = std::make_shared<T>(torch::div(a31, det_));
      result[3] = std::make_shared<T>(torch::div(a41, det_));
      result[4] = std::make_shared<T>(torch::div(a12, det_));
      result[5] = std::make_shared<T>(torch::div(a22, det_));
      result[6] = std::make_shared<T>(torch::div(a32, det_));
      result[7] = std::make_shared<T>(torch::div(a42, det_));
      result[8] = std::make_shared<T>(torch::div(a13, det_));
      result[9] = std::make_shared<T>(torch::div(a23, det_));
      result[10] = std::make_shared<T>(torch::div(a33, det_));
      result[11] = std::make_shared<T>(torch::div(a43, det_));
      result[12] = std::make_shared<T>(torch::div(a14, det_));
      result[13] = std::make_shared<T>(torch::div(a24, det_));
      result[14] = std::make_shared<T>(torch::div(a34, det_));
      result[15] = std::make_shared<T>(torch::div(a44, det_));
      return result;
    } else {
      throw std::runtime_error("Unsupported block tensor dimension");
      return *this;
    }
  }
 
  inline auto ginvtr() const {
    if constexpr (Rows == Cols)
      return this->invtr();
    else
      // Compute the transpose of the generalized inverse, i.e. A (A^T A)^{-T}
      return (*this) * (this->tr() * (*this)).invtr();
  }
 
  inline auto trace() const {
    static_assert(Rows == Cols, "trace(.) requires square block tensor");
 
    if constexpr (Rows == 1)
      return BlockTensor<T, 1, 1>(*Base::data_[0]);
 
    else if constexpr (Rows == 2)
      return BlockTensor<T, 1, 1>(*Base::data_[0] + *Base::data_[3]);
 
    else if constexpr (Rows == 3)
      return BlockTensor<T, 1, 1>(*Base::data_[0] + *Base::data_[4] +
                                  *Base::data_[8]);
 
    else if constexpr (Rows == 4)
      return BlockTensor<T, 1, 1>(*Base::data_[0] + *Base::data_[5] +
                                  *Base::data_[10] + *Base::data_[15]);
 
    else
      throw std::runtime_error("Unsupported block tensor dimension");
  }
 
private:
  template <std::size_t... Is>
  inline auto norm_(std::index_sequence<Is...>) const {
    return torch::sqrt(std::apply([](const auto&... tensors) {
      return (tensors + ...);
    }, std::make_tuple(std::get<Is>(Base::data_)->square()...)));
  }
  
public:
  inline auto norm() const {
    return BlockTensor<T, 1, 1>(std::make_shared<T>(norm_(std::make_index_sequence<Rows*Cols>{})));
  }
  
private:
  template <std::size_t... Is>
  inline auto normalize_(std::index_sequence<Is...> is) const {
    auto n_ = norm_(is);
    return BlockTensor<T, Rows, Cols>(std::make_shared<T>(*std::get<Is>(Base::data_)/n_)...);
  }
  
public:
  inline auto normalize() const {
    return normalize_(std::make_index_sequence<Rows*Cols>{});
  }
 
private:
  template <std::size_t... Is>
  inline auto dot_(std::index_sequence<Is...>, const BlockTensor<T, Rows, Cols> &other) const {
    return std::apply([](const auto&... tensors) {
      return (tensors + ...);
    }, std::make_tuple(torch::mul(*std::get<Is>(Base::data_), *std::get<Is>(other.data_))...));    
  }
  
public:
  inline auto dot(const BlockTensor<T, Rows, Cols> & other) const {
    return BlockTensor<T, 1, 1>(std::make_shared<T>(dot_(std::make_index_sequence<Rows*Cols>{}, other)));
  }
 
  inline virtual void
  pretty_print(std::ostream &os = Log(log::info)) const noexcept override {
    os << Base::name() << "\n";
    for (std::size_t row = 0; row < Rows; ++row)
      for (std::size_t col = 0; col < Cols; ++col)
        os << "[" << row << "," << col << "] = \n"
           << *Base::data_[Cols * row + col] << "\n";
  }
};
  
template <typename T, typename U, std::size_t Rows, std::size_t Common,
          std::size_t Cols>
inline auto operator*(const BlockTensor<T, Rows, Common> &lhs,
                      const BlockTensor<U, Common, Cols> &rhs) {
  BlockTensor<typename std::common_type<T, U>::type, Rows, Cols> result;
  for (std::size_t row = 0; row < Rows; ++row)
    for (std::size_t col = 0; col < Cols; ++col) {
      T tmp =
          (lhs[Common * row]->dim() > rhs[col]->dim()
               ? torch::mul(*lhs[Common * row], rhs[col]->unsqueeze(-1))
               : (lhs[Common * row]->dim() < rhs[col]->dim()
                      ? torch::mul(lhs[Common * row]->unsqueeze(-1), *rhs[col])
                      : torch::mul(*lhs[Common * row], *rhs[col])));
      for (std::size_t idx = 1; idx < Common; ++idx)
        tmp += (lhs[Common * row]->dim() > rhs[col]->dim()
                    ? torch::mul(*lhs[Common * row + idx],
                                 rhs[Cols * idx + col]->unsqueeze(-1))
                    : (lhs[Common * row]->dim() < rhs[col]->dim()
                           ? torch::mul(lhs[Common * row + idx]->unsqueeze(-1),
                                        *rhs[Cols * idx + col])
                           : torch::mul(*lhs[Common * row + idx],
                                        *rhs[Cols * idx + col])));
      result[Cols * row + col] = std::make_shared<T>(tmp);
    }
  return result;
}
 
template <typename T, std::size_t Rows, std::size_t Cols, std::size_t Slices>
class BlockTensor<T, Rows, Cols, Slices>
    : public BlockTensorCore<T, Rows, Cols, Slices> {
private:
  using Base = BlockTensorCore<T, Rows, Cols, Slices>;
 
public:
  using BlockTensorCore<T, Rows, Cols, Slices>::BlockTensorCore;
 
  inline static constexpr std::size_t rows() { return Rows; }
 
  inline static constexpr std::size_t cols() { return Cols; }
 
  inline static constexpr std::size_t slices() { return Slices; }
 
  using Base::operator();
 
  inline const T &operator()(std::size_t row, std::size_t col,
                             std::size_t slice) const {
    assert(0 <= row && row < Rows && 0 <= col && col < Cols && 0 <= slice &&
           slice < Slices);
    return *Base::data_[Rows * Cols * slice + Cols * row + col];
  }
 
  inline T &operator()(std::size_t row, std::size_t col, std::size_t slice) {
    assert(0 <= row && row < Rows && 0 <= col && col < Cols && 0 <= slice &&
           slice < Slices);
    return *Base::data_[Rows * Cols * slice + Cols * row + col];
  }
 
  using Base::set;
 
  template <typename D>
  inline T &set(std::size_t row, std::size_t col, std::size_t slice, D &&data) {
    Base::data_[Rows * Cols * slice + Cols * row + col] =
        make_shared<D>(std::forward<D>(data));
    return *Base::data_[Rows * Cols * slice + Cols * row + col];
  }
 
  inline auto slice(std::size_t slice) const {
    assert(0 <= slice && slice < Slices);
    BlockTensor<T, Rows, Cols> result;
    for (std::size_t row = 0; row < Rows; ++row)
      for (std::size_t col = 0; col < Cols; ++col)
        result[Cols * row + col] =
            Base::data_[Rows * Cols * slice + Cols * row + col];
    return result;
  }
 
  inline auto reorder_ikj() const {
    BlockTensor<T, Rows, Slices, Cols> result;
    for (std::size_t slice = 0; slice < Slices; ++slice)
      for (std::size_t row = 0; row < Rows; ++row)
        for (std::size_t col = 0; col < Cols; ++col)
          result[Rows * Slices * col + Slices * row + slice] =
              Base::data_[Rows * Cols * slice + Cols * row + col];
    return result;
  }
 
  inline auto reorder_jik() const {
    BlockTensor<T, Cols, Rows, Slices> result;
    for (std::size_t slice = 0; slice < Slices; ++slice)
      for (std::size_t row = 0; row < Rows; ++row)
        for (std::size_t col = 0; col < Cols; ++col)
          result[Rows * Cols * slice + Rows * col + row] =
              Base::data_[Rows * Cols * slice + Cols * row + col];
    return result;
  }
 
  inline auto reorder_kji() const {
    BlockTensor<T, Slices, Cols, Rows> result;
    for (std::size_t slice = 0; slice < Slices; ++slice)
      for (std::size_t row = 0; row < Rows; ++row)
        for (std::size_t col = 0; col < Cols; ++col)
          result[Slices * Cols * row + Cols * slice + col] =
              Base::data_[Rows * Cols * slice + Cols * row + col];
    return result;
  }
 
  inline auto reorder_kij() const {
    BlockTensor<T, Slices, Rows, Cols> result;
    for (std::size_t slice = 0; slice < Slices; ++slice)
      for (std::size_t row = 0; row < Rows; ++row)
        for (std::size_t col = 0; col < Cols; ++col)
          result[Slices * Rows * col + Rows * slice + row] =
              Base::data_[Rows * Cols * slice + Cols * row + col];
    return result;
  }
 
  inline virtual void
  pretty_print(std::ostream &os = Log(log::info)) const noexcept override {
    os << Base::name() << "\n";
    for (std::size_t slice = 0; slice < Slices; ++slice)
      for (std::size_t row = 0; row < Rows; ++row)
        for (std::size_t col = 0; col < Cols; ++col)
          os << "[" << row << "," << col << "," << slice << "] = \n"
             << *Base::data_[Rows * Cols * slice + Cols * row + col] << "\n";
  }
};
 
template <typename T, typename U, std::size_t Rows, std::size_t Common,
          std::size_t Cols, std::size_t Slices>
inline auto operator*(const BlockTensor<T, Rows, Common> &lhs,
                      const BlockTensor<U, Common, Cols, Slices> &rhs) {
  BlockTensor<typename std::common_type<T, U>::type, Rows, Cols, Slices> result;
  for (std::size_t slice = 0; slice < Slices; ++slice)
    for (std::size_t row = 0; row < Rows; ++row)
      for (std::size_t col = 0; col < Cols; ++col) {
        T tmp =
            (lhs[Common * row]->dim() > rhs[Rows * Cols * slice + col]->dim()
                 ? torch::mul(*lhs[Common * row],
                              rhs[Rows * Cols * slice + col]->unsqueeze(-1))
                 : (lhs[Common * row]->dim() <
                            rhs[Rows * Cols * slice + col]->dim()
                        ? torch::mul(lhs[Common * row]->unsqueeze(-1),
                                     *rhs[Rows * Cols * slice + col])
                        : torch::mul(*lhs[Common * row],
                                     *rhs[Rows * Cols * slice + col])));
        for (std::size_t idx = 1; idx < Common; ++idx)
          tmp +=
              (lhs[Common * row]->dim() > rhs[Rows * Cols * slice + col]->dim()
                   ? torch::mul(
                         *lhs[Common * row + idx],
                         rhs[Rows * Cols * slice + Cols * idx + col]->unsqueeze(
                             -1))
                   : (lhs[Common * row]->dim() <
                              rhs[Rows * Cols * slice + col]->dim()
                          ? torch::mul(
                                lhs[Common * row + idx]->unsqueeze(-1),
                                *rhs[Rows * Cols * slice + Cols * idx + col])
                          : torch::mul(
                                *lhs[Common * row + idx],
                                *rhs[Rows * Cols * slice + Cols * idx + col])));
        result[Rows * Cols * slice + Cols * row + col] =
            std::make_shared<T>(tmp);
      }
  return result;
}
 
template <typename T, typename U, std::size_t Rows, std::size_t Common,
          std::size_t Cols, std::size_t Slices>
inline auto operator*(const BlockTensor<T, Rows, Common, Slices> &lhs,
                      const BlockTensor<U, Common, Cols> &rhs) {
  BlockTensor<typename std::common_type<T, U>::type, Rows, Cols, Slices> result;
  for (std::size_t slice = 0; slice < Slices; ++slice)
    for (std::size_t row = 0; row < Rows; ++row)
      for (std::size_t col = 0; col < Cols; ++col) {
        T tmp =
            (lhs[Rows * Cols * slice + Common * row]->dim() > rhs[col]->dim()
                 ? torch::mul(*lhs[Rows * Cols * slice + Common * row],
                              rhs[col]->unsqueeze(-1))
                 : (lhs[Rows * Cols * slice + Common * row]->dim() <
                            rhs[col]->dim()
                        ? torch::mul(lhs[Rows * Cols * slice + Common * row]
                                         ->unsqueeze(-1),
                                     *rhs[col])
                        : torch::mul(*lhs[Rows * Cols * slice + Common * row],
                                     *rhs[col])));
        for (std::size_t idx = 1; idx < Common; ++idx)
          tmp +=
              (lhs[Rows * Cols * slice + Common * row + idx]->dim() >
                       rhs[Cols * idx + col]->dim()
                   ? torch::mul(*lhs[Rows * Cols * slice + Common * row + idx],
                                rhs[Cols * idx + col])
                         ->unsqueeze(-1)
                   : (lhs[Rows * Cols * slice + Common * row + idx]->dim() <
                              rhs[Cols * idx + col]->dim()
                          ? torch::mul(
                                lhs[Rows * Cols * slice + Common * row + idx]
                                    ->unsqueeze(-1),
                                *rhs[Cols * idx + col])
                          : torch::mul(
                                *lhs[Rows * Cols * slice + Common * row + idx],
                                *rhs[Cols * idx + col])));
        result[Rows * Cols * slice + Cols * row + col] =
            std::make_shared<T>(tmp);
      }
  return result;
}
 
#define blocktensor_unary_op(name)                                             \
  template <typename T, std::size_t... Dims>                                   \
  inline auto name(const BlockTensor<T, Dims...> &input) {                     \
    BlockTensor<T, Dims...> result;                                            \
    for (std::size_t idx = 0; idx < (Dims * ...); ++idx)                       \
      result[idx] = std::make_shared<T>(torch::name(*input[idx]));             \
    return result;                                                             \
  }
 
#define blocktensor_unary_special_op(name)                                     \
  template <typename T, std::size_t... Dims>                                   \
  inline auto name(const BlockTensor<T, Dims...> &input) {                     \
    BlockTensor<T, Dims...> result;                                            \
    for (std::size_t idx = 0; idx < (Dims * ...); ++idx)                       \
      result[idx] = std::make_shared<T>(torch::special::name(*input[idx]));    \
    return result;                                                             \
  }
 
#define blocktensor_binary_op(name)                                            \
  template <typename T, typename U, std::size_t... Dims>                       \
  inline auto name(const BlockTensor<T, Dims...> &input,                       \
                   const BlockTensor<U, Dims...> &other) {                     \
    BlockTensor<typename std::common_type<T, U>::type, Dims...> result;        \
    for (std::size_t idx = 0; idx < (Dims * ...); ++idx)                       \
      result[idx] =                                                            \
          std::make_shared<T>(torch::name(*input[idx], *other[idx]));          \
    return result;                                                             \
  }
 
#define blocktensor_binary_special_op(name)                                    \
  template <typename T, typename U, std::size_t... Dims>                       \
  inline auto name(const BlockTensor<T, Dims...> &input,                       \
                   const BlockTensor<U, Dims...> &other) {                     \
    BlockTensor<typename std::common_type<T, U>::type, Dims...> result;        \
    for (std::size_t idx = 0; idx < (Dims * ...); ++idx)                       \
      result[idx] =                                                            \
          std::make_shared<T>(torch::special::name(*input[idx], *other[idx])); \
    return result;                                                             \
  }
 
blocktensor_unary_op(abs);
 
blocktensor_unary_op(absolute);
 
blocktensor_unary_op(acos);
 
blocktensor_unary_op(arccos);
 
blocktensor_unary_op(acosh);
 
blocktensor_unary_op(arccosh);
 
template <typename T, typename U, typename V, std::size_t... Dims>
inline auto add(const BlockTensor<T, Dims...> &input,
                const BlockTensor<U, Dims...> &other, V alpha = 1.0) {
  BlockTensor<typename std::common_type<T, U>::type, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] =
        std::make_shared<T>(torch::add(*input[idx], *other[idx], alpha));
  return result;
}
 
template <typename T, typename U, typename V, std::size_t... Dims>
inline auto add(const BlockTensor<T, Dims...> &input, U other, V alpha = 1.0) {
  BlockTensor<T, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(torch::add(*input[idx], other, alpha));
  return result;
}
 
template <typename T, typename U, typename V, std::size_t... Dims>
inline auto add(T input, const BlockTensor<U, Dims...> &other, V alpha = 1.0) {
  BlockTensor<U, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(torch::add(input, *other[idx], alpha));
  return result;
}
 
template <typename T, typename U, typename V, typename W, std::size_t... Dims>
inline auto addcdiv(const BlockTensor<T, Dims...> &input,
                    const BlockTensor<U, Dims...> &tensor1,
                    const BlockTensor<V, Dims...> &tensor2, W value = 1.0) {
  BlockTensor<typename std::common_type<T, U, V>::type, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(
        torch::addcdiv(*input[idx], *tensor1[idx], *tensor2[idx], value));
  return result;
}
 
template <typename T, typename U, typename V, typename W, std::size_t... Dims>
inline auto addcmul(const BlockTensor<T, Dims...> &input,
                    const BlockTensor<U, Dims...> &tensor1,
                    const BlockTensor<V, Dims...> &tensor2, W value = 1.0) {
  BlockTensor<typename std::common_type<T, U, V>::type, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(
        torch::addcmul(*input[idx], *tensor1[idx], *tensor2[idx], value));
  return result;
}
 
blocktensor_unary_op(angle);
 
blocktensor_unary_op(asin);
 
blocktensor_unary_op(arcsin);
 
blocktensor_unary_op(asinh);
 
blocktensor_unary_op(arcsinh);
 
blocktensor_unary_op(atan);
 
blocktensor_unary_op(arctan);
 
blocktensor_unary_op(atanh)
 
    
    blocktensor_unary_op(arctanh);
 
blocktensor_binary_op(atan2);
 
#if TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11 ||                   \
    TORCH_VERSION_MAJOR >= 2
blocktensor_binary_op(arctan2);
#endif
 
blocktensor_unary_op(bitwise_not);
 
blocktensor_binary_op(bitwise_and);
 
blocktensor_binary_op(bitwise_or);
 
blocktensor_binary_op(bitwise_xor);
 
blocktensor_binary_op(bitwise_left_shift);
 
blocktensor_binary_op(bitwise_right_shift);
 
blocktensor_unary_op(ceil);
 
template <typename T, typename U, std::size_t... Dims>
inline auto clamp(const BlockTensor<T, Dims...> &input, U min, U max) {
  BlockTensor<T, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(torch::clamp(*input[idx], min, max));
  return result;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto clip(const BlockTensor<T, Dims...> &input, U min, U max) {
  BlockTensor<T, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(torch::clip(*input[idx], min, max));
  return result;
}
 
blocktensor_unary_op(conj_physical);
 
blocktensor_binary_op(copysign);
 
blocktensor_unary_op(cos);
 
blocktensor_unary_op(cosh);
 
blocktensor_unary_op(deg2rad)
 
    
    blocktensor_binary_op(div);
 
blocktensor_binary_op(divide);
 
blocktensor_unary_op(digamma);
 
template <typename T, std::size_t Rows, std::size_t Cols>
inline auto dot(const BlockTensor<T, Rows, Cols> &input, const BlockTensor<T, Rows, Cols> &tensor) {
  return input.dot(tensor);
}
  
blocktensor_unary_op(erf);
 
blocktensor_unary_op(erfc);
 
blocktensor_unary_op(erfinv);
 
blocktensor_unary_op(exp);
 
blocktensor_unary_op(exp2);
 
blocktensor_unary_op(expm1);
 
blocktensor_unary_op(fix);
 
blocktensor_binary_op(float_power);
 
blocktensor_unary_op(floor);
 
blocktensor_binary_op(fmod);
 
blocktensor_unary_op(frac);
 
blocktensor_unary_op(frexp);
 
blocktensor_unary_op(imag);
 
blocktensor_binary_op(ldexp);
 
blocktensor_unary_op(lgamma);
 
blocktensor_unary_op(log);
 
blocktensor_unary_op(log10);
 
blocktensor_unary_op(log1p);
 
blocktensor_unary_op(log2);
 
blocktensor_binary_op(logaddexp);
 
blocktensor_binary_op(logaddexp2);
 
blocktensor_binary_op(logical_and)
 
    
    blocktensor_unary_op(logical_not)
 
    blocktensor_binary_op(logical_or)
 
    blocktensor_binary_op(logical_xor);
 
 
blocktensor_binary_op(hypot);
 
blocktensor_unary_op(i0);
 
blocktensor_binary_special_op(gammainc);
 
blocktensor_binary_op(igamma);
 
blocktensor_binary_special_op(gammaincc);
 
blocktensor_binary_op(igammac);
 
blocktensor_binary_op(mul);
 
blocktensor_binary_op(multiply);
 
blocktensor_unary_op(neg);
 
blocktensor_unary_op(negative);
 
blocktensor_binary_op(nextafter);
 
blocktensor_unary_op(positive);
 
blocktensor_binary_op(pow);
 
blocktensor_unary_op(rad2deg);
 
blocktensor_unary_op(real);
 
blocktensor_unary_op(reciprocal);
 
blocktensor_binary_op(remainder);
 
blocktensor_unary_op(round);
 
blocktensor_unary_op(rsqrt);
 
blocktensor_unary_special_op(expit);
 
blocktensor_unary_op(sigmoid);
 
blocktensor_unary_op(sign);
 
blocktensor_unary_op(sgn);
 
blocktensor_unary_op(signbit);
 
blocktensor_unary_op(sin);
 
blocktensor_unary_op(sinc);
 
blocktensor_unary_op(sinh);
 
blocktensor_unary_op(sqrt);
 
blocktensor_unary_op(square);
 
template <typename T, typename U, typename V, std::size_t... Dims>
inline auto sub(const BlockTensor<T, Dims...> &input,
                const BlockTensor<U, Dims...> &other, V alpha = 1.0) {
  BlockTensor<typename std::common_type<T, U>::type, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] =
        std::make_shared<T>(torch::sub(*input[idx], *other[idx], alpha));
  return result;
}
 
template <typename T, typename U, typename V, std::size_t... Dims>
inline auto subtract(const BlockTensor<T, Dims...> &input,
                     const BlockTensor<U, Dims...> &other, V alpha = 1.0) {
  BlockTensor<typename std::common_type<T, U>::type, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] =
        std::make_shared<T>(torch::sub(*input[idx], *other[idx], alpha));
  return result;
}
 
blocktensor_unary_op(tan);
 
blocktensor_unary_op(tanh);
 
blocktensor_unary_op(trunc)
 
    
    blocktensor_binary_op(xlogy);
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator+(const BlockTensor<T, Dims...> &lhs,
                      const BlockTensor<U, Dims...> &rhs) {
  BlockTensor<typename std::common_type<T, U>::type, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(*lhs[idx] + *rhs[idx]);
  return result;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator+(const BlockTensor<T, Dims...> &lhs, const U &rhs) {
  BlockTensor<T, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(*lhs[idx] + rhs);
  return result;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator+(const T &lhs, const BlockTensor<U, Dims...> &rhs) {
  BlockTensor<U, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<U>(lhs + *rhs[idx]);
  return result;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator+=(BlockTensor<T, Dims...> &lhs,
                       const BlockTensor<U, Dims...> &rhs) {
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    lhs[idx] = std::make_shared<T>(*lhs[idx] + *rhs[idx]);
  return lhs;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator+=(BlockTensor<T, Dims...> &lhs, const U &rhs) {
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    lhs[idx] = std::make_shared<T>(*lhs[idx] + rhs);
  return lhs;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator-(const BlockTensor<T, Dims...> &lhs,
                      const BlockTensor<U, Dims...> &rhs) {
  BlockTensor<typename std::common_type<T, U>::type, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(*lhs[idx] - *rhs[idx]);
  return result;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator-(const BlockTensor<T, Dims...> &lhs, const U &rhs) {
  BlockTensor<T, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<T>(*lhs[idx] - rhs);
  return result;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator-(const T &lhs, const BlockTensor<U, Dims...> &rhs) {
  BlockTensor<U, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] = std::make_shared<U>(lhs - *rhs[idx]);
  return result;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator-=(BlockTensor<T, Dims...> &lhs,
                       const BlockTensor<U, Dims...> &rhs) {
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    lhs[idx] = std::make_shared<T>(*lhs[idx] - *rhs[idx]);
  return lhs;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator-=(BlockTensor<T, Dims...> &lhs, const U &rhs) {
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    lhs[idx] = std::make_shared<T>(*lhs[idx] - rhs);
  return lhs;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator*(const BlockTensor<T, Dims...> &lhs, const U &rhs) {
  BlockTensor<T, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] =
        (lhs[idx]->dim() > rhs.dim()
             ? std::make_shared<T>(*lhs[idx] * rhs.unsqueeze(-1))
             : (lhs[idx]->dim() < rhs.dim()
                    ? std::make_shared<T>(lhs[idx]->unsqueeze(-1) * rhs)
                    : std::make_shared<T>(*lhs[idx] * rhs)));
  ;
  return result;
}
 
template <typename T, typename U, std::size_t... Dims>
inline auto operator*(const T &lhs, const BlockTensor<U, Dims...> &rhs) {
  BlockTensor<U, Dims...> result;
  for (std::size_t idx = 0; idx < (Dims * ...); ++idx)
    result[idx] =
        (lhs.dim() > rhs[idx]->dim()
             ? std::make_shared<U>(lhs * rhs[idx]->unsqueeze(-1))
             : (lhs.dim() < rhs[idx]->dim()
                    ? std::make_shared<U>(lhs.unsqueeze(-1) * *rhs[idx])
                    : std::make_shared<U>(lhs * *rhs[idx])));
  return result;
}
 
template <typename T, typename U, std::size_t... TDims, std::size_t... UDims>
inline bool operator==(const BlockTensor<T, TDims...> &lhs,
                       const BlockTensor<U, UDims...> &rhs) {
  if constexpr ((sizeof...(TDims) != sizeof...(UDims)) ||
                ((TDims != UDims) || ...))
    return false;
 
  bool result = true;
  for (std::size_t idx = 0; idx < (TDims * ...); ++idx)
    result = result && torch::equal(*lhs[idx], *rhs[idx]);
 
  return result;
}
 
template <typename T, typename U, std::size_t... TDims, std::size_t... UDims>
inline bool operator!=(const BlockTensor<T, TDims...> &lhs,
                       const BlockTensor<U, UDims...> &rhs) {
  return !(lhs == rhs);
}
 
} // namespace utils
} // namespace iganet