tensor_impl.hpp

Go to the documentation of this file.

/*
  mptensor - Parallel Library for Tensor Network Methods
 
  Copyright 2016 Satoshi Morita
 
  mptensor is free software: you can redistribute it and/or modify it
  under the terms of the GNU Lesser General Public License as
  published by the Free Software Foundation, either version 3 of the
  License, or (at your option) any later version.
 
  mptensor is distributed in the hope that it will be useful, but
  WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
 
  You should have received a copy of the GNU Lesser General Public
  License along with mptensor.  If not, see
  <https://www.gnu.org/licenses/>.
*/
 
#ifndef _TENSOR_IMPL_HPP_
#define _TENSOR_IMPL_HPP_
 
#include <algorithm>
#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <type_traits>
#include <vector>
 
#include "mptensor/complex.hpp"
#include "mptensor/index.hpp"
#include "mptensor/matrix/matrix.hpp"
#include "mptensor/tensor.hpp"
 
namespace mptensor {
 
/* Utilities */
bool is_no_transpose(const Axes &axes, const Axes &axes_map, size_t rank);
namespace debug {
bool check_total_size(const Shape &s1, const Shape &s2);
bool check_extend(const Shape &s_old, const Shape &s_new);
bool check_transpose_axes(const Axes &axes, size_t rank);
bool check_svd_axes(const Axes &axes_row, const Axes &axes_col, size_t rank);
bool check_trace_axes(const Axes &axes_1, const Axes &axes_2, size_t rank);
bool check_trace_axes(const Axes &axes_a, const Axes &axes_b,
                      const Shape &shape_a, const Shape &shape_b);
bool check_contract_axes(const Axes &axes_1, const Axes &axes_2, size_t rank);
bool check_square(const Shape &shape, size_t urank);
}  // namespace debug
 
/* ---------- constructors ---------- */
 
 
template <typename MatrixType>
Tensor<MatrixType>::Tensor() : Mat(), upper_rank(0), axes_map(){};
 
 
template <typename MatrixType>
Tensor<MatrixType>::Tensor(const Shape &shape) : Mat() {
  init(shape, shape.size() / 2);
};
 
 
template <typename MatrixType>
Tensor<MatrixType>::Tensor(const comm_type &comm)
    : Mat(comm), upper_rank(0), axes_map(){};
 
 
template <typename MatrixType>
Tensor<MatrixType>::Tensor(const comm_type &comm, const Shape &shape)
    : Mat(comm) {
  init(shape, shape.size() / 2);
};
 
 
template <typename MatrixType>
Tensor<MatrixType>::Tensor(const comm_type &comm, const Shape &shape,
                          const size_t upper_rank)
    : Mat(comm) {
  init(shape, upper_rank);
};
 
 
template <typename MatrixType>
Tensor<MatrixType>::Tensor(const comm_type &comm,
                          const Tensor<lapack::Matrix<value_type>> &t)
    : Mat(comm) {
  init(t.shape(), t.get_upper_rank());
  const size_t n = Mat.local_size();
  size_t idx;
  int dummy;
  for (size_t i = 0; i < n; ++i) {
    t.local_position(global_index(i), dummy, idx);
    Mat[i] = t[idx];
  }
};
 
 
template <typename MatrixType>
Tensor<MatrixType>::Tensor(const comm_type &comm, const std::vector<value_type> &v)
    : Mat(comm) {
  init(Shape(v.size()), 0);
  const size_t n = Mat.local_size();
  Index idx;
  for (size_t i = 0; i < n; ++i) {
    idx = global_index(i);
    Mat[i] = v[idx[0]];
  }
};
 
/* ---------- inline member functions ---------- */
 
template <typename MatrixType>
inline const Shape &Tensor<MatrixType>::shape() const {
  return Dim;
};
 
template <typename MatrixType>
inline size_t Tensor<MatrixType>::rank() const {
  return Dim.size();
};
 
 
template <typename MatrixType>
inline size_t Tensor<MatrixType>::ndim() const {
  return Dim.size();
};
 
template <typename MatrixType>
inline size_t Tensor<MatrixType>::local_size() const {
  return Mat.local_size();
};
 
template <typename MatrixType>
inline size_t Tensor<MatrixType>::get_upper_rank() const {
  return upper_rank;
};
 
 
template <typename MatrixType>
inline const Axes &Tensor<MatrixType>::get_axes_map() const {
  return axes_map;
};
 
template <typename MatrixType>
inline const MatrixType &Tensor<MatrixType>::get_matrix() const {
  return Mat;
};
 
template <typename MatrixType>
inline MatrixType &Tensor<MatrixType>::get_matrix() {
  return Mat;
};
 
template <typename MatrixType>
inline const typename Tensor<MatrixType>::comm_type &
Tensor<MatrixType>::get_comm() const {
  return Mat.get_comm();
};
 
template <typename MatrixType>
inline int Tensor<MatrixType>::get_comm_size() const {
  return Mat.get_comm_size();
};
 
template <typename MatrixType>
inline int Tensor<MatrixType>::get_comm_rank() const {
  return Mat.get_comm_rank();
};
 
 
template <typename MatrixType>
inline const typename MatrixType::value_type &Tensor<MatrixType>::operator[](size_t local_idx) const {
  return Mat[local_idx];
}
 
 
template <typename MatrixType>
inline typename MatrixType::value_type &Tensor<MatrixType>::operator[](size_t local_idx) {
  return Mat[local_idx];
}
 
template <typename MatrixType>
inline void Tensor<MatrixType>::prep_global_to_local() const {
  Mat.prep_global_to_local();
};
 
template <typename MatrixType>
inline void Tensor<MatrixType>::prep_local_to_global() const {
  Mat.prep_local_to_global();
};
 
template <typename MatrixType>
inline void Tensor<MatrixType>::make_l2g_map() const {
  const size_t rank = Dim.size();
  const size_t rank0 = upper_rank;
  const size_t rank1 = rank - rank0;  // lower_rank
  const size_t l_row = Mat.local_row_size();
  const size_t l_col = Mat.local_col_size();
  // When rank0 (== upper_rank) or rank1 (== lower_rank) is zero, the
  // corresponding range is empty; taking the base address of an empty range
  // (&axes_map[rank0] one past the end, or &l2g_map_row[0] on an empty vector)
  // is undefined behavior and aborts under _GLIBCXX_ASSERTIONS. Guard with a
  // null pointer that is never dereferenced (the matching loops are skipped).
  const size_t *axes_map0 = rank0 > 0 ? &(axes_map[0]) : nullptr;
  const size_t *axes_map1 = rank1 > 0 ? &(axes_map[rank0]) : nullptr;
 
  l2g_map_row.resize(l_row * rank0);
  l2g_map_col.resize(l_col * rank1);
 
#pragma omp parallel default(shared)
  {
    // Type int instead of size_t is used because of std::div()
    std::vector<int> dim0(rank0);
    std::vector<int> dim1(rank1);
    for (size_t i = 0; i < rank0; ++i) dim0[i] = int(Dim[axes_map0[i]]);
    for (size_t i = 0; i < rank1; ++i) dim1[i] = int(Dim[axes_map1[i]]);
 
    int g_row, g_col;  // global matrix row- and column-index
    std::div_t divresult;
    size_t *map;
    if (rank0 > 0) {
#pragma omp for
      for (size_t k = 0; k < l_row; ++k) {
        map = &(l2g_map_row[k * rank0]);
        g_row = int(Mat.global_row_index(k));
        for (size_t i = 0; i < rank0; ++i) {
          divresult = std::div(g_row, dim0[i]);
          map[i] = divresult.rem;
          g_row = divresult.quot;
        }
      }
    }
    if (rank1 > 0) {
#pragma omp for
      for (size_t k = 0; k < l_col; ++k) {
        map = &(l2g_map_col[k * rank1]);
        g_col = int(Mat.global_col_index(k));
        for (size_t i = 0; i < rank1; ++i) {
          divresult = std::div(g_col, dim1[i]);
          map[i] = divresult.rem;
          g_col = divresult.quot;
        }
      }
    }
  }
 
  return;
};
 
 
template <typename MatrixType>
inline void Tensor<MatrixType>::global_index_l2g_map(size_t lindex,
                                                    size_t gindex[]) const {
  const size_t rank = Dim.size();
  const size_t rank0 = upper_rank;
  const size_t rank1 = rank - rank0;  // lower_rank
  const size_t lindex_row = Mat.local_row_index(lindex);
  const size_t lindex_col = Mat.local_col_index(lindex);
  // Guard against empty upper/lower ranges (see make_l2g_map): a null base
  // pointer is safe because the matching loop below is then empty.
  const size_t *map_row = rank0 > 0 ? &(l2g_map_row[lindex_row * rank0]) : nullptr;
  const size_t *map_col = rank1 > 0 ? &(l2g_map_col[lindex_col * rank1]) : nullptr;
  const size_t *axes_map0 = rank0 > 0 ? &(axes_map[0]) : nullptr;
  const size_t *axes_map1 = rank1 > 0 ? &(axes_map[rank0]) : nullptr;
  for (size_t i = 0; i < rank0; ++i) {
    gindex[axes_map0[i]] = map_row[i];
  }
  for (size_t i = 0; i < rank1; ++i) {
    gindex[axes_map1[i]] = map_col[i];
  }
  return;
};
 
 
template <typename MatrixType>
inline void Tensor<MatrixType>::global_index_l2g_map_transpose(
    size_t lindex, const size_t axes_trans[], size_t index_new[]) const {
  const size_t rank = Dim.size();
  const size_t rank0 = upper_rank;
  const size_t rank1 = rank - rank0;  // lower_rank
  const size_t lindex_row = Mat.local_row_index(lindex);
  const size_t lindex_col = Mat.local_col_index(lindex);
  // Guard against empty upper/lower ranges (see make_l2g_map): a null base
  // pointer is safe because the matching loop below is then empty.
  const size_t *map_row = rank0 > 0 ? &(l2g_map_row[lindex_row * rank0]) : nullptr;
  const size_t *map_col = rank1 > 0 ? &(l2g_map_col[lindex_col * rank1]) : nullptr;
  // const size_t* axes_map0 = &(axes_map[0]);
  // const size_t* axes_map1 = &(axes_map[rank0]);
  for (size_t i = 0; i < rank0; ++i) {
    index_new[axes_trans[i]] = map_row[i];
    // index_new[axes_inv[axes_map0[i]]] = map_row[i];
  }
  for (size_t i = 0; i < rank1; ++i) {
    index_new[axes_trans[i + rank0]] = map_col[i];
    // index_new[axes_inv[axes_map1[i]]] = map_col[i];
  }
  return;
};
 
template <typename MatrixType>
inline void Tensor<MatrixType>::local_position_fast(size_t g_row, size_t g_col,
                                                   int &comm_rank,
                                                   size_t &local_idx) const {
  Mat.local_position(g_row, g_col, comm_rank, local_idx);
}
 
/* ---------- member functions ---------- */
 
 
template <typename MatrixType>
inline void Tensor<MatrixType>::init(const Shape &shape, size_t urank) {
  init(shape, urank, range(shape.size()));
}
 
 
template <typename MatrixType>
void Tensor<MatrixType>::init(const Shape &shape, size_t urank,
                             const Axes &map) {
  Dim = shape;
  axes_map = map;
  const size_t rank = Dim.size();
  assert(urank <= rank);
  upper_rank = urank;
 
  size_t n_row = 1;
  size_t n_col = 1;
  for (size_t i = 0; i < upper_rank; ++i) n_row *= shape[map[i]];
  for (size_t i = upper_rank; i < rank; ++i) n_col *= shape[map[i]];
  Mat.init(n_row, n_col);
}
 
 
template <typename MatrixType>
void Tensor<MatrixType>::print_info(std::ostream &out,
                                   const std::string &tag) const {
  if (tag.size() > 0) {
    out << tag << ":\t";
  }
  out << "Tensor: shape= " << Dim << " upper_rank= " << upper_rank
      << " axes_map= " << axes_map << "\t";
  Mat.print_info(out);
};
 
 
template <typename MatrixType>
void Tensor<MatrixType>::print_info_mpi(std::ostream &out,
                                       const std::string &tag) const {
  const int mpisize = get_comm_size();
  const int mpirank = get_comm_rank();
  const typename Tensor<MatrixType>::comm_type &comm = get_comm();
  Mat.barrier();
  for (int i = 0; i < mpisize; ++i) {
    if (i == mpirank) {
      if (tag.size() > 0) {
        out << tag << ":\t";
      }
      out << "mpirank: " << mpirank << "\t";
      print_info(out, "");
    }
    Mat.barrier();
  }
}
 
 
template <typename MatrixType>
bool Tensor<MatrixType>::local_index(const Index &gindex, size_t &lindex) const {
  const size_t rank = gindex.size();
  assert(rank == Dim.size());
  assert(rank > 0);
  size_t g_row(0), g_col(0);
  size_t d_row(1), d_col(1);
  for (size_t i = 0; i < upper_rank; ++i) {
    size_t j = axes_map[i];
    g_row += gindex[j] * d_row;
    d_row *= Dim[j];
  }
  for (size_t i = upper_rank; i < rank; ++i) {
    size_t j = axes_map[i];
    g_col += gindex[j] * d_col;
    d_col *= Dim[j];
  }
  return Mat.local_index(g_row, g_col, lindex);
};
 
 
template <typename MatrixType>
Index Tensor<MatrixType>::global_index(size_t lindex) const {
  const size_t rank = Dim.size();
  Index gindex;
  size_t g_row, g_col;
  Mat.global_index(lindex, g_row, g_col);
  gindex.resize(rank);
  std::div_t divresult;
  for (size_t i = 0; i < upper_rank; ++i) {
    size_t j = axes_map[i];
    divresult = std::div(int(g_row), int(Dim[j]));
    gindex[j] = divresult.rem;
    g_row = divresult.quot;
  }
  for (size_t i = upper_rank; i < rank; ++i) {
    size_t j = axes_map[i];
    divresult = std::div(int(g_col), int(Dim[j]));
    gindex[j] = divresult.rem;
    g_col = divresult.quot;
  }
  return gindex;
};
 
 
template <typename MatrixType>
void Tensor<MatrixType>::global_index_fast(size_t lindex, Index &gindex) const {
  const size_t rank = Dim.size();
  size_t g_row, g_col;
  Mat.global_index(lindex, g_row, g_col);
  std::div_t divresult;
  for (size_t i = 0; i < upper_rank; ++i) {
    size_t j = axes_map[i];
    divresult = std::div(int(g_row), int(Dim[j]));
    gindex[j] = divresult.rem;
    g_row = divresult.quot;
  }
  for (size_t i = upper_rank; i < rank; ++i) {
    size_t j = axes_map[i];
    divresult = std::div(int(g_col), int(Dim[j]));
    gindex[j] = divresult.rem;
    g_col = divresult.quot;
  }
  return;
};
 
 
template <typename MatrixType>
bool Tensor<MatrixType>::get_value(const Index &idx, value_type &val) const {
  size_t li;
  if (local_index(idx, li)) {
    val = Mat[li];
    return true;
  } else {
    return false;
  }
}
 
 
template <typename MatrixType>
void Tensor<MatrixType>::set_value(const Index &idx, value_type val) {
  size_t li;
  if (local_index(idx, li)) {
    Mat[li] = val;
  }
}
 
 
template <typename MatrixType>
void Tensor<MatrixType>::local_position(const Index &index, int &comm_rank,
                                       size_t &local_idx) const {
  const size_t rank = Dim.size();
  size_t g_row(0), g_col(0);
  size_t d_row(1), d_col(1);
  for (size_t i = 0; i < upper_rank; ++i) {
    size_t j = axes_map[i];
    g_row += index[j] * d_row;
    d_row *= Dim[j];
  }
  for (size_t i = upper_rank; i < rank; ++i) {
    size_t j = axes_map[i];
    g_col += index[j] * d_col;
    d_col *= Dim[j];
  }
  Mat.local_position(g_row, g_col, comm_rank, local_idx);
  return;
}
 
 
template <typename MatrixType>
void Tensor<MatrixType>::change_configuration(const size_t new_upper_rank,
                                             const Axes &new_axes_map) {
  if ((upper_rank == new_upper_rank) && (axes_map == new_axes_map)) return;
 
  Tensor<MatrixType> T_old(*this);
 
  const size_t rank = this->rank();
  Shape dim;
  Axes axes;
  dim.resize(rank);
  axes.resize(rank);
  for (size_t i = 0; i < rank; ++i) {
    dim[i] = Dim[new_axes_map[i]];
    axes[new_axes_map[i]] = i;
  }
  init(dim, new_upper_rank);
  transpose(axes);
 
  /* create lists of local position and destination rank */
  const size_t local_size = T_old.local_size();
  std::vector<int> dest_mpirank(local_size);
  std::vector<size_t> local_position(local_size);
 
  T_old.prep_local_to_global();
  this->prep_global_to_local();
 
#pragma omp parallel default(shared)
  {
    Index index;
    index.resize(rank);
 
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      T_old.global_index_fast(i, index);
 
      int dest;
      size_t pos;
      this->local_position(index, dest, pos);
 
      local_position[i] = pos;
      dest_mpirank[i] = dest;
    }
  }
 
  /* exchange data */
  replace_matrix_data(T_old.get_matrix(), dest_mpirank, local_position, Mat);
 
  return;
}
 
 
template <typename MatrixType>
Tensor<MatrixType> &Tensor<MatrixType>::transpose(const Axes &axes) {
  const size_t rank = Dim.size();
  assert(debug::check_transpose_axes(axes, rank));
 
  Shape dim_now = Dim;
  Axes map_now = axes_map;
  Axes axes_inv;
  axes_inv.resize(rank);
  for (size_t i = 0; i < rank; ++i) {
    axes_inv[axes[i]] = i;
  }
 
  for (size_t i = 0; i < rank; ++i) {
    Dim[i] = dim_now[axes[i]];
    axes_map[i] = axes_inv[map_now[i]];
  }
 
  return (*this);
}
 
 
template <typename MatrixType>
template <typename D>
Tensor<MatrixType> &Tensor<MatrixType>::multiply_vector(const std::vector<D> &vec,
                                                      size_t n_axes) {
  assert(Dim[n_axes] <= vec.size());
  const size_t local_size = this->local_size();
  prep_local_to_global();
#pragma omp parallel default(shared)
  {
    Index idx;
    idx.resize(rank());
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      global_index_fast(i, idx);
      Mat[i] *= vec[idx[n_axes]];
    }
  }
  return (*this);
};
 
 
template <typename MatrixType>
template <typename D0, typename D1>
Tensor<MatrixType> &Tensor<MatrixType>::multiply_vector(
    const std::vector<D0> &vec0, size_t n_axes0, const std::vector<D1> &vec1,
    size_t n_axes1) {
  assert(Dim[n_axes0] <= vec0.size());
  assert(Dim[n_axes1] <= vec1.size());
  const size_t local_size = this->local_size();
  prep_local_to_global();
#pragma omp parallel default(shared)
  {
    Index idx;
    idx.resize(rank());
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      global_index_fast(i, idx);
      Mat[i] *= vec0[idx[n_axes0]] * vec1[idx[n_axes1]];
    }
  }
  return (*this);
};
 
 
template <typename MatrixType>
template <typename D0, typename D1, typename D2>
Tensor<MatrixType> &Tensor<MatrixType>::multiply_vector(
    const std::vector<D0> &vec0, size_t n_axes0, const std::vector<D1> &vec1,
    size_t n_axes1, const std::vector<D2> &vec2, size_t n_axes2) {
  assert(Dim[n_axes0] <= vec0.size());
  assert(Dim[n_axes1] <= vec1.size());
  assert(Dim[n_axes2] <= vec2.size());
  const size_t local_size = this->local_size();
  prep_local_to_global();
#pragma omp parallel default(shared)
  {
    Index idx;
    idx.resize(rank());
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      global_index_fast(i, idx);
      Mat[i] *= vec0[idx[n_axes0]] * vec1[idx[n_axes1]] * vec2[idx[n_axes2]];
    }
  }
  return (*this);
};
 
 
template <typename MatrixType>
template <typename D0, typename D1, typename D2, typename D3>
Tensor<MatrixType> &Tensor<MatrixType>::multiply_vector(
    const std::vector<D0> &vec0, size_t n_axes0, const std::vector<D1> &vec1,
    size_t n_axes1, const std::vector<D2> &vec2, size_t n_axes2,
    const std::vector<D3> &vec3, size_t n_axes3) {
  assert(Dim[n_axes0] <= vec0.size());
  assert(Dim[n_axes1] <= vec1.size());
  assert(Dim[n_axes2] <= vec2.size());
  assert(Dim[n_axes3] <= vec3.size());
  const size_t local_size = this->local_size();
  prep_local_to_global();
#pragma omp parallel default(shared)
  {
    Index idx;
    idx.resize(rank());
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      global_index_fast(i, idx);
      Mat[i] *= vec0[idx[n_axes0]] * vec1[idx[n_axes1]] * vec2[idx[n_axes2]] *
                vec3[idx[n_axes3]];
    }
  }
  return (*this);
};
 
 
template <typename MatrixType>
Tensor<MatrixType> &Tensor<MatrixType>::set_slice(const Tensor<MatrixType> &a,
                                                const size_t n_axes,
                                                const size_t i_begin,
                                                const size_t i_end) {
  assert(rank() == a.rank());
  assert(n_axes < rank());
  assert(i_begin < i_end);
  assert(i_end <= shape()[n_axes]);
  assert(i_end - i_begin == a.shape()[n_axes]);
 
  /* create lists of local position and destination rank */
  const size_t local_size = a.local_size();
  std::vector<int> dest_mpirank(local_size);
  std::vector<size_t> local_position(local_size);
 
  a.prep_local_to_global();
  prep_global_to_local();
 
#pragma omp parallel default(shared)
  {
    Index index;
    index.resize(rank());
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      a.global_index_fast(i, index);
      index[n_axes] += i_begin;
 
      int dest;
      size_t pos;
      this->local_position(index, dest, pos);
 
      local_position[i] = pos;
      dest_mpirank[i] = dest;
    }
  }
 
  /* exchange data */
  replace_matrix_data(a.get_matrix(), dest_mpirank, local_position,
                      get_matrix());
 
  return (*this);
}
 
 
template <typename MatrixType>
Tensor<MatrixType> &Tensor<MatrixType>::set_slice(const Tensor<MatrixType> &a,
                                                const Index &index_begin,
                                                const Index &index_end) {
  const size_t nr = rank();
  assert(nr == a.rank());
  assert(index_begin.size() == nr);
  assert(index_end.size() == nr);
  for (size_t r = 0; r < nr; ++r) {
    assert(index_end[r] <= shape()[r]);
    assert(index_begin[r] <= index_end[r]);
  }
 
  /* create lists of local position and destination rank */
  const size_t local_size = a.local_size();
  std::vector<int> dest_mpirank(local_size);
  std::vector<size_t> local_position(local_size);
 
  a.prep_local_to_global();
  prep_global_to_local();
 
#pragma omp parallel default(shared)
  {
    Index index;
    index.resize(nr);
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      a.global_index_fast(i, index);
      for (size_t r = 0; r < nr; ++r) {
        if (index_begin[r] != index_end[r]) index[r] += index_begin[r];
      }
 
      int dest;
      size_t pos;
      this->local_position(index, dest, pos);
 
      local_position[i] = pos;
      dest_mpirank[i] = dest;
    }
  }
 
  /* exchange data */
  replace_matrix_data(a.get_matrix(), dest_mpirank, local_position,
                      get_matrix());
 
  return (*this);
}
 
 
 
template <typename MatrixType>
Tensor<lapack::Matrix<typename MatrixType::value_type>> Tensor<MatrixType>::gather() {
  const size_t n = rank();
  if (!(axes_map == range(n))) {
    change_configuration(n / 2, range(n));
  }
  Tensor<lapack::Matrix<typename MatrixType::value_type>> T(get_comm(), get_matrix().flatten());
  return reshape(T, Dim);
}
 
 
 
template <typename MatrixType>
std::vector<typename MatrixType::value_type> Tensor<MatrixType>::flatten() {
  const size_t n = rank();
  if (!(axes_map == range(n))) {
    change_configuration(n / 2, range(n));
  }
  return get_matrix().flatten();
}
 
template <typename MatrixType>
Tensor<MatrixType> &Tensor<MatrixType>::operator+=(const Tensor<MatrixType> &rhs) {
  assert(Dim == rhs.shape());
  change_configuration(rhs.get_upper_rank(), rhs.get_axes_map());
  Mat += rhs.get_matrix();
  return (*this);
}
 
template <typename MatrixType>
Tensor<MatrixType> &Tensor<MatrixType>::operator-=(const Tensor<MatrixType> &rhs) {
  assert(Dim == rhs.shape());
  change_configuration(rhs.get_upper_rank(), rhs.get_axes_map());
  Mat -= rhs.get_matrix();
  return (*this);
}
 
template <typename MatrixType>
Tensor<MatrixType> &Tensor<MatrixType>::operator*=(value_type rhs) {
  Mat *= rhs;
  return (*this);
}
 
template <typename MatrixType>
Tensor<MatrixType> &Tensor<MatrixType>::operator/=(value_type rhs) {
  Mat /= rhs;
  return (*this);
}
 
template <typename MatrixType>
Tensor<MatrixType> &Tensor<MatrixType>::operator=(value_type rhs) {
  const size_t n = Mat.local_size();
  for (int i = 0; i < n; ++i) Mat[i] = rhs;
  return (*this);
}
 
 
template <typename MatrixType>
template <typename UnaryOperation>
Tensor<MatrixType> &Tensor<MatrixType>::map(UnaryOperation op) {
  Mat.map(op);
  return (*this);
}
 
/* ---------- non-member functions (Opeations) ---------- */
 
 
template <typename MatrixType>
inline Tensor<MatrixType> transpose(Tensor<MatrixType> T, const Axes &axes) {
  return T.transpose(axes);
}
 
 
template <typename MatrixType>
Tensor<MatrixType> transpose(const Tensor<MatrixType> &T, const Axes &axes,
                            size_t urank_new) {
  const size_t rank = T.rank();
  assert(debug::check_transpose_axes(axes, rank));
  assert(urank_new <= rank);
 
  if (urank_new == T.get_upper_rank()) {
    if (is_no_transpose(axes, T.get_axes_map(), rank)) return T;
  }
 
  /* new index dimension */
  Shape dim_old = T.shape();
  Shape dim_new;
  dim_new.resize(rank);
  for (size_t r = 0; r < rank; ++r) {
    dim_new[r] = dim_old[axes[r]];
  }
 
  /* new tensor */
  Tensor<MatrixType> T_new(T.get_comm(), dim_new, urank_new);
 
  /* create lists of local position and destination rank */
  const size_t local_size = T.local_size();
  std::vector<int> dest_mpirank(local_size);
  std::vector<size_t> local_position(local_size);
 
  T.make_l2g_map();
  T_new.prep_global_to_local();
 
#pragma omp parallel default(shared)
  {
    // Assuming axes_map of T_new is [0,1,2,...]
    std::vector<size_t> d_offset(rank);
    size_t d_row(1), d_col(1);
    for (size_t r = 0; r < urank_new; ++r) {
      d_offset[r] = d_row;
      d_row *= dim_new[r];
    }
    for (size_t r = urank_new; r < rank; ++r) {
      d_offset[r] = d_col;
      d_col *= dim_new[r];
    }
 
    Axes axes_map = T.get_axes_map();
    std::vector<size_t> axes_inv(rank);
    std::vector<size_t> axes_trans(rank);
    for (size_t r = 0; r < rank; ++r) {
      axes_inv[axes[r]] = r;
    }
    for (size_t r = 0; r < rank; ++r) {
      axes_trans[r] = axes_inv[axes_map[r]];
    }
 
    std::vector<size_t> index_new(rank);
 
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      T.global_index_l2g_map_transpose(i, axes_trans.data(), index_new.data());
 
      size_t g_row(0),
          g_col(0);  // indices of global matrix for transposed tensor
      for (size_t r = 0; r < urank_new; ++r) {
        g_row += index_new[r] * d_offset[r];
      }
      for (size_t r = urank_new; r < rank; ++r) {
        g_col += index_new[r] * d_offset[r];
      }
 
      T_new.local_position_fast(g_row, g_col, dest_mpirank[i],
                                local_position[i]);
    }
  }
 
  /* exchange data */
  replace_matrix_data(T.get_matrix(), dest_mpirank, local_position,
                      T_new.get_matrix());
 
  return T_new;
};
 
 
template <typename MatrixType>
Tensor<MatrixType> reshape(const Tensor<MatrixType> &T, const Shape &shape_new) {
  assert(debug::check_total_size(shape_new, T.shape()));
 
  const Shape &shape = T.shape();
  const size_t rank = shape.size();
  const size_t rank_new = shape_new.size();
 
  /* initialize new tensor */
  Tensor<MatrixType> T_new(T.get_comm(), shape_new);
 
  /* create lists of local position and destination rank */
  const size_t local_size = T.local_size();
  std::vector<int> dest_mpirank(local_size);
  std::vector<size_t> local_position(local_size);
 
  T.prep_local_to_global();
  T_new.prep_global_to_local();
 
#pragma omp parallel default(shared)
  {
    Index index, index_new;
    index.resize(rank);
    index_new.resize(rank_new);
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      T.global_index_fast(i, index);
      size_t global_position(0);
      size_t dim(1);
      for (size_t r = 0; r < rank; ++r) {
        global_position += index[r] * dim;
        dim *= shape[r];
      }
 
      for (size_t r = 0; r < rank_new; ++r) {
        index_new[r] = global_position % shape_new[r];
        global_position /= shape_new[r];
      }
 
      int dest;
      size_t pos;
      T_new.local_position(index_new, dest, pos);
 
      local_position[i] = pos;
      dest_mpirank[i] = dest;
    }
  }
 
  /* exchange data */
  replace_matrix_data(T.get_matrix(), dest_mpirank, local_position,
                      T_new.get_matrix());
 
  return T_new;
};
 
 
template <typename MatrixType>
Tensor<MatrixType> slice(const Tensor<MatrixType> &T, size_t n_axes,
                        size_t i_begin, size_t i_end) {
  const int mpisize = T.get_comm_size();
  const Shape &shape = T.shape();
  assert(n_axes < T.rank());
  assert(i_begin < i_end);
  assert(i_end <= shape[n_axes]);
 
  Shape shape_new = shape;
  shape_new[n_axes] = i_end - i_begin;
 
  /* initialize new tensor */
  Tensor<MatrixType> T_new(T.get_comm(), shape_new);
 
  /* create lists of local position and destination rank */
  const size_t local_size = T.local_size();
  std::vector<int> dest_mpirank(local_size);
  std::vector<size_t> local_position(local_size);
 
  T.prep_local_to_global();
  T_new.prep_global_to_local();
 
#pragma omp parallel default(shared)
  {
    Index index;
    index.resize(T.rank());
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      T.global_index_fast(i, index);
      if (index[n_axes] >= i_begin && index[n_axes] < i_end) {
        index[n_axes] -= i_begin;
 
        int dest;
        size_t pos;
        T_new.local_position(index, dest, pos);
 
        local_position[i] = pos;
        dest_mpirank[i] = dest;
      } else {
        // not send
        local_position[i] = 0;
        dest_mpirank[i] = mpisize;
      }
    }
  }
 
  /* exchange data */
  replace_matrix_data(T.get_matrix(), dest_mpirank, local_position,
                      T_new.get_matrix());
 
  return T_new;
}
 
 
template <typename MatrixType>
Tensor<MatrixType> slice(const Tensor<MatrixType> &T, const Index &index_begin,
                        const Index &index_end) {
  const int mpisize = T.get_comm_size();
  const Shape &shape = T.shape();
  const size_t rank = T.rank();
  assert(index_begin.size() == rank);
  assert(index_end.size() == rank);
 
  Shape shape_new = shape;
  for (size_t r = 0; r < rank; ++r) {
    assert(index_end[r] <= shape[r]);
    assert(index_begin[r] <= index_end[r]);
    shape_new[r] = index_end[r] - index_begin[r];
    if (shape_new[r] == 0) shape_new[r] = shape[r];
  }
 
  /* initialize new tensor */
  Tensor<MatrixType> T_new(T.get_comm(), shape_new);
 
  /* create lists of local position and destination rank */
  const size_t local_size = T.local_size();
  std::vector<int> dest_mpirank(local_size);
  std::vector<size_t> local_position(local_size);
 
  T.prep_local_to_global();
  T_new.prep_global_to_local();
 
#pragma omp parallel default(shared)
  {
    Index index;
    index.resize(T.rank());
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      T.global_index_fast(i, index);
 
      bool is_send = true;
      for (size_t r = 0; r < rank; ++r) {
        size_t idx = index[r];
        if (index_begin[r] == index_end[r]) {
          continue;
        } else if (idx >= index_begin[r] && idx < index_end[r]) {
          index[r] -= index_begin[r];
        } else {
          is_send = false;
          break;
        }
      }
 
      if (is_send) {
        int dest;
        size_t pos;
        T_new.local_position(index, dest, pos);
 
        local_position[i] = pos;
        dest_mpirank[i] = dest;
      } else {
        // not send
        local_position[i] = 0;
        dest_mpirank[i] = mpisize;
      }
    }
  }
 
  /* exchange data */
  replace_matrix_data(T.get_matrix(), dest_mpirank, local_position,
                      T_new.get_matrix());
 
  return T_new;
};
 
 
template <typename MatrixType>
Tensor<MatrixType> extend(const Tensor<MatrixType> &T, const Shape &shape_new) {
  assert(T.rank() == shape_new.size());
  assert(debug::check_extend(T.shape(), shape_new));
 
  /* initialize new tensor */
  Tensor<MatrixType> T_new(T.get_comm(), shape_new);
 
  /* create lists of local position and destination rank */
  const size_t local_size = T.local_size();
  std::vector<int> dest_mpirank(local_size);
  std::vector<size_t> local_position(local_size);
 
  T.prep_local_to_global();
  T_new.prep_global_to_local();
 
#pragma omp parallel default(shared)
  {
    Index index;
    index.resize(T.rank());
#pragma omp for
    for (size_t i = 0; i < local_size; ++i) {
      T.global_index_fast(i, index);
 
      int dest;
      size_t pos;
      T_new.local_position(index, dest, pos);
 
      local_position[i] = pos;
      dest_mpirank[i] = dest;
    }
  }
 
  /* exchange data */
  replace_matrix_data(T.get_matrix(), dest_mpirank, local_position,
                      T_new.get_matrix());
 
  return T_new;
}
 
 
template <typename MatrixType>
typename MatrixType::value_type trace(const Tensor<MatrixType> &M) {
  assert(M.rank() == 2);
  return matrix_trace(M.get_matrix());
};
 
 
template <typename MatrixType>
typename MatrixType::value_type trace(const Tensor<MatrixType> &T, const Axes &axes_1, const Axes &axes_2) {
  assert(axes_1.size() == axes_2.size());
  assert(axes_1.size() + axes_2.size() == T.rank());
  assert(debug::check_trace_axes(axes_1, axes_2, T.rank()));
 
  if (T.rank() == 2) return trace(T);
 
  const size_t n = T.local_size();
  const size_t l = axes_1.size();
  Index index;
  bool check;
  typename MatrixType::value_type sum(0.0);
 
  index.resize(T.rank());
  for (size_t i = 0; i < n; ++i) {
    T.global_index_fast(i, index);
    check = true;
    for (size_t k = 0; k < l; ++k) {
      if (index[axes_1[k]] != index[axes_2[k]]) {
        check = false;
        break;
      }
    }
    if (check) {
      sum += T[i];
    }
  }
 
  return T.get_matrix().allreduce_sum(sum);
};
 
 
template <typename MatrixType>
typename MatrixType::value_type trace(const Tensor<MatrixType> &A, const Tensor<MatrixType> &B,
        const Axes &axes_a, const Axes &axes_b) {
  assert(A.rank() == B.rank());
  assert(A.rank() == axes_a.size());
  assert(B.rank() == axes_b.size());
  assert(debug::check_trace_axes(axes_a, axes_b, A.shape(), B.shape()));
 
  const size_t rank = A.rank();
  Axes axes;
  Axes axes_a_inv;
  Axes axes_map = A.get_axes_map();
  axes.resize(rank);
  axes_a_inv.resize(rank);
 
  for (size_t i = 0; i < rank; ++i) {
    axes_a_inv[axes_a[i]] = i;
  }
  for (size_t i = 0; i < rank; ++i) {
    axes[i] = axes_b[axes_a_inv[axes_map[i]]];
  }
 
  const size_t n = A.local_size();
  Tensor<MatrixType> B_t = transpose(B, axes, A.get_upper_rank());
  typename MatrixType::value_type sum(0.0);
 
  for (size_t i = 0; i < n; ++i) {
    sum += A[i] * B_t[i];
  }
 
  return A.get_matrix().allreduce_sum(sum);
}
 
 
template <typename MatrixType>
Tensor<MatrixType> contract(const Tensor<MatrixType> &T, const Axes &axes_1,
                           const Axes &axes_2) {
  const int mpisize = T.get_comm_size();
  assert(axes_1.size() == axes_2.size());
  assert(axes_1.size() + axes_2.size() < T.rank());
  if (axes_1.size() == 0) return T;
  assert(debug::check_contract_axes(axes_1, axes_2, T.rank()));
 
  Shape shape = T.shape();
  Shape shape_new;
  Axes axes_new;
  {
    Axes v = axes_1 + axes_2;
    v.sort();
    size_t k = 0;
    size_t n = v.size();
    for (size_t i = 0; i < T.rank(); ++i) {
      if (k < n && i == v[k]) {
        ++k;
      } else {
        shape_new.push(shape[i]);
        axes_new.push(i);
      }
    }
  }
 
  /* initialize new tensor */
  Tensor<MatrixType> T_new(T.get_comm(), shape_new);
 
  /* create lists of local position and destination rank */
  std::vector<int> dest_mpirank(T.local_size());
  std::vector<size_t> local_position(T.local_size());
 
  Index index, index_new;
  index.resize(T.rank());
  index_new.resize(T_new.rank());
  bool check;
  const size_t n = axes_1.size();
 
  for (size_t i = 0; i < T.local_size(); ++i) {
    T.global_index_fast(i, index);
    check = true;
    for (size_t k = 0; k < n; ++k) {
      if (index[axes_1[k]] != index[axes_2[k]]) {
        check = false;
        continue;
      }
    }
 
    if (check) {
      for (size_t k = 0; k < index_new.size(); ++k) {
        index_new[k] = index[axes_new[k]];
      }
 
      int dest;
      size_t pos;
      T_new.local_position(index_new, dest, pos);
 
      local_position[i] = pos;
      dest_mpirank[i] = dest;
    } else {
      // not send
      local_position[i] = 0;
      dest_mpirank[i] = mpisize;
    }
  }
 
  /* sum data */
  sum_matrix_data(T.get_matrix(), dest_mpirank, local_position,
                  T_new.get_matrix());
 
  return T_new;
}
 
 
template <typename MatrixType>
Tensor<MatrixType> kron(const Tensor<MatrixType> &a, const Tensor<MatrixType> &b) {
  assert(a.rank() == b.rank());
  assert(a.get_comm() == b.get_comm());
 
  Shape shape_a = a.shape();
  Shape shape_b = b.shape();
  Shape shape_c = shape_a;
  size_t n = shape_a.size();
  Axes axes_trans;
  axes_trans.resize(2 * n);
  for (size_t i = 0; i < shape_b.size(); ++i) {
    shape_c[i] *= shape_b[i];
    axes_trans[2 * i] = i;
    axes_trans[2 * i + 1] = i + n;
  }
 
  return reshape(
      transpose(tensordot(reshape(a, shape_a + Shape(1)),
                          reshape(b, Shape(1) + shape_b), Axes(n), Axes(0)),
                axes_trans),
      shape_c);
};
 
 
template <typename MatrixType>
Tensor<MatrixType> tensordot(const Tensor<MatrixType> &a,
                            const Tensor<MatrixType> &b, const Axes &axes_a,
                            const Axes &axes_b) {
  assert(axes_a.size() == axes_b.size());
  assert(a.get_comm() == b.get_comm());
  Shape shape_a = a.shape();
  Shape shape_b = b.shape();
  for (size_t i = 0; i < axes_a.size(); ++i) {
    assert(shape_a[axes_a[i]] == shape_b[axes_b[i]]);
  }
 
  const typename Tensor<MatrixType>::comm_type &comm = a.get_comm();
 
  Shape shape_c;
  const size_t rank_row_c = a.rank() - axes_a.size();
  const size_t rank_col_c = b.rank() - axes_b.size();
  shape_c.resize(rank_row_c + rank_col_c);
 
  Axes trans_axes_a;
  Axes trans_axes_b;
  size_t urank_a;
  size_t urank_b;
 
  {
    const size_t rank = a.rank();
    const size_t rank_row = rank - axes_a.size();
    const size_t rank_col = axes_a.size();
    std::vector<size_t> v(rank);
    for (size_t i = 0; i < rank; ++i) v[i] = i;
    for (size_t i = 0; i < rank_col; ++i) v[axes_a[i]] = rank;
    std::sort(v.begin(), v.end());
    for (size_t i = 0; i < rank_col; ++i) v[rank_row + i] = axes_a[i];
 
    trans_axes_a.assign(rank, v.data());
    urank_a = rank_row;
 
    for (size_t i = 0; i < rank_row; ++i) shape_c[i] = shape_a[v[i]];
  }
 
  {
    const size_t rank = b.rank();
    const size_t rank_row = axes_b.size();
    const size_t rank_col = rank - axes_b.size();
    std::vector<size_t> v(rank);
    for (size_t i = 0; i < rank; ++i) v[i] = i;
    for (size_t i = 0; i < rank_row; ++i) v[axes_b[i]] = 0;
    std::sort(v.begin(), v.end());
    for (size_t i = 0; i < rank_row; ++i) v[i] = axes_b[i];
 
    trans_axes_b.assign(rank, v.data());
    urank_b = rank_row;
 
    for (size_t i = 0; i < rank_col; ++i)
      shape_c[i + rank_row_c] = shape_b[v[i + rank_row]];
  }
 
  Tensor<MatrixType> c(comm, shape_c, rank_row_c);
  matrix_product(transpose(a, trans_axes_a, urank_a).get_matrix(),
                 transpose(b, trans_axes_b, urank_b).get_matrix(),
                 c.get_matrix());
 
  return c;
};
 
 
template <typename MatrixType>
int svd(const Tensor<MatrixType> &a, std::vector<double> &s) {
  assert(a.rank() == 2);
  int info;
  info = svd(a, Axes(0), Axes(1), s);
  return info;
}
 
 
template <typename MatrixType>
int svd(const Tensor<MatrixType> &a, Tensor<MatrixType> &u,
        std::vector<double> &s, Tensor<MatrixType> &vt) {
  assert(a.rank() == 2);
  int info;
  info = svd(a, Axes(0), Axes(1), u, s, vt);
  return info;
}
 
 
template <typename MatrixType>
int svd(const Tensor<MatrixType> &a, const Axes &axes_row, const Axes &axes_col,
        std::vector<double> &s) {
  assert(axes_row.size() > 0);
  assert(axes_col.size() > 0);
  assert(debug::check_svd_axes(axes_row, axes_col, a.rank()));
 
  const size_t rank = a.rank();
  const size_t urank = axes_row.size();
 
  Axes axes = axes_row + axes_col;
 
  Tensor<MatrixType> a_t = transpose(a, axes, urank);
  const Shape &shape = a_t.shape();
 
  size_t d_row(1), d_col(1);
  for (size_t i = 0; i < urank; ++i) d_row *= shape[i];
  for (size_t i = urank; i < rank; ++i) d_col *= shape[i];
  size_t size = (d_row < d_col) ? d_row : d_col;
 
  s.resize(size);
 
  int info;
  info = matrix_svd(a_t.get_matrix(), s);
 
  return info;
}
 
 
template <typename MatrixType>
int svd(const Tensor<MatrixType> &a, const Axes &axes_row, const Axes &axes_col,
        Tensor<MatrixType> &u, std::vector<double> &s, Tensor<MatrixType> &vt) {
  assert(axes_row.size() > 0);
  assert(axes_col.size() > 0);
  assert(debug::check_svd_axes(axes_row, axes_col, a.rank()));
 
  const size_t rank = a.rank();
  const size_t urank = axes_row.size();
 
  Axes axes = axes_row + axes_col;
 
  Tensor<MatrixType> a_t = transpose(a, axes, urank);
  const Shape &shape = a_t.shape();
 
  size_t d_row(1), d_col(1);
  for (size_t i = 0; i < urank; ++i) d_row *= shape[i];
  for (size_t i = urank; i < rank; ++i) d_col *= shape[i];
  size_t size = (d_row < d_col) ? d_row : d_col;
 
  Shape shape_u;
  shape_u.resize(urank + 1);
  for (size_t i = 0; i < urank; ++i) shape_u[i] = shape[i];
  shape_u[urank] = size;
 
  Shape shape_vt;
  shape_vt.resize(rank - urank + 1);
  shape_vt[0] = size;
  for (size_t i = urank; i < rank; ++i) shape_vt[i - urank + 1] = shape[i];
 
  size_t urank_u = urank;
  size_t urank_vt = 1;
 
  u = Tensor<MatrixType>(a.get_comm(), shape_u, urank_u);
  vt = Tensor<MatrixType>(a.get_comm(), shape_vt, urank_vt);
  s.resize(size);
 
  int info;
  info = matrix_svd(a_t.get_matrix(), u.get_matrix(), s, vt.get_matrix());
 
  return info;
}
 
 
template <typename MatrixType>
int psvd(const Tensor<MatrixType> &a, std::vector<double> &s,
         const size_t target_rank) {
  int info = svd(a, s);
  if (s.size() > target_rank) s.resize(target_rank);
  return info;
}
 
 
template <typename MatrixType>
int psvd(const Tensor<MatrixType> &a, Tensor<MatrixType> &u,
         std::vector<double> &s, Tensor<MatrixType> &vt,
         const size_t target_rank) {
  int info = svd(a, u, s, vt);
  if (s.size() > target_rank) {
    s.resize(target_rank);
    u = slice(u, 1, 0, target_rank);
    vt = slice(vt, 0, 0, target_rank);
  }
  return info;
}
 
 
template <typename MatrixType>
int psvd(const Tensor<MatrixType> &a, const Axes &axes_row, const Axes &axes_col,
         std::vector<double> &s, const size_t target_rank) {
  int info = svd(a, axes_row, axes_col, s);
  if (s.size() > target_rank) s.resize(target_rank);
  return info;
}
 
 
template <typename MatrixType>
int psvd(const Tensor<MatrixType> &a, const Axes &axes_row, const Axes &axes_col,
         Tensor<MatrixType> &u, std::vector<double> &s, Tensor<MatrixType> &vt,
         const size_t target_rank) {
  int info = svd(a, axes_row, axes_col, u, s, vt);
  if (s.size() > target_rank) {
    s.resize(target_rank);
    u = slice(u, axes_row.size(), 0, target_rank);
    vt = slice(vt, 0, 0, target_rank);
  }
  return info;
}
 
 
template <typename MatrixType>
int qr(const Tensor<MatrixType> &a, Tensor<MatrixType> &q, Tensor<MatrixType> &r) {
  assert(a.rank() == 2);
  int info;
  info = qr(a, Axes(0), Axes(1), q, r);
  return info;
}
 
 
template <typename MatrixType>
int qr(const Tensor<MatrixType> &a, const Axes &axes_row, const Axes &axes_col,
       Tensor<MatrixType> &q, Tensor<MatrixType> &r) {
  assert(axes_row.size() > 0);
  assert(axes_col.size() > 0);
  assert(debug::check_svd_axes(axes_row, axes_col, a.rank()));
 
  const size_t rank = a.rank();
  const size_t urank = axes_row.size();
  const Shape shape_a = a.shape();
 
  Axes axes = axes_row + axes_col;
  Shape shape;
  shape.resize(rank);
  for (size_t i = 0; i < rank; ++i) shape[i] = shape_a[axes[i]];
 
  size_t d_row(1), d_col(1);
  for (size_t i = 0; i < urank; ++i) d_row *= shape[i];
  for (size_t i = urank; i < rank; ++i) d_col *= shape[i];
  size_t size = (d_row < d_col) ? d_row : d_col;
 
  // rank-2 tensors (matrices)
  Tensor<MatrixType> mat_q =
      reshape(transpose(a, axes, urank), Shape(d_row, d_col));
  Tensor<MatrixType> mat_r(mat_q.get_comm(), mat_q.shape(), 1);
 
  // QR decomposition. Elements of mat_q change from a to q.
  int info;
  info = matrix_qr(mat_q.get_matrix(), mat_r.get_matrix());
 
  // Get shape of q and r.
  Shape shape_q;
  shape_q.resize(urank + 1);
  for (size_t i = 0; i < urank; ++i) shape_q[i] = shape[i];
  shape_q[urank] = size;
 
  Shape shape_r;
  shape_r.resize(rank - urank + 1);
  shape_r[0] = size;
  for (size_t i = urank; i < rank; ++i) shape_r[i - urank + 1] = shape[i];
 
  // Reshape q and r.
  if (d_row > d_col) {
    q = reshape(mat_q, shape_q);
    r = reshape(slice(mat_r, 0, 0, size), shape_r);
  } else if (d_row < d_col) {
    q = reshape(slice(mat_q, 1, 0, size), shape_q);
    r = reshape(mat_r, shape_r);
  } else {  // d_row==d_col
    q = reshape(mat_q, shape_q);
    r = reshape(mat_r, shape_r);
  }
 
  return info;
}
 
 
template <typename MatrixType>
int eigh(const Tensor<MatrixType> &a, std::vector<double> &w,
         Tensor<MatrixType> &z) {
  assert(a.rank() == 2);
  Shape shape = a.shape();
  assert(shape[0] == shape[1]);
 
  Tensor<MatrixType> a_t = transpose(a, Axes(0, 1), 1);
 
  size_t n = shape[0];
  w.resize(n);
  z = Tensor<MatrixType>(a.get_comm(), Shape(n, n), 1);
 
  int info;
  info = matrix_eigh(a_t.get_matrix(), w, z.get_matrix());
  return info;
};
 
 
template <typename MatrixType>
int eigh(const Tensor<MatrixType> &a, std::vector<double> &w) {
  assert(a.rank() == 2);
  Shape shape = a.shape();
  assert(shape[0] == shape[1]);
 
  Tensor<MatrixType> a_t = transpose(a, Axes(0, 1), 1);
 
  size_t n = shape[0];
  w.resize(n);
 
  int info;
  info = matrix_eigh(a_t.get_matrix(), w);
  return info;
};
 
 
template <typename MatrixType>
int eigh(const Tensor<MatrixType> &a, const Axes &axes_row, const Axes &axes_col,
         std::vector<double> &w, Tensor<MatrixType> &z) {
  assert(axes_row.size() > 0);
  assert(axes_col.size() > 0);
 
  const size_t rank = a.rank();
  const size_t urank = axes_row.size();
 
  Axes axes = axes_row + axes_col;
 
  Tensor<MatrixType> a_t = transpose(a, axes, urank);
  const Shape &shape = a_t.shape();
 
  size_t d_row(1), d_col(1);
  for (size_t i = 0; i < urank; ++i) d_row *= shape[i];
  for (size_t i = urank; i < rank; ++i) d_col *= shape[i];
  size_t size = (d_row < d_col) ? d_row : d_col;
 
  assert(d_row == d_col);
 
  Shape shape_z;
  shape_z.resize(urank + 1);
  for (size_t i = 0; i < urank; ++i) shape_z[i] = shape[i];
  shape_z[urank] = size;
 
  z = Tensor<MatrixType>(a.get_comm(), shape_z, urank);
  w.resize(size);
 
  int info;
  info = matrix_eigh(a_t.get_matrix(), w, z.get_matrix());
  return info;
};
 
 
template <typename MatrixType>
int eigh(const Tensor<MatrixType> &a, const Axes &axes_row, const Axes &axes_col,
         std::vector<double> &w) {
  assert(axes_row.size() > 0);
  assert(axes_col.size() > 0);
 
  const size_t rank = a.rank();
  const size_t urank = axes_row.size();
 
  Axes axes = axes_row + axes_col;
 
  Tensor<MatrixType> a_t = transpose(a, axes, urank);
  const Shape &shape = a_t.shape();
 
  size_t d_row(1), d_col(1);
  for (size_t i = 0; i < urank; ++i) d_row *= shape[i];
  for (size_t i = urank; i < rank; ++i) d_col *= shape[i];
  size_t size = (d_row < d_col) ? d_row : d_col;
  w.resize(size);
 
  assert(d_row == d_col);
 
  int info;
  info = matrix_eigh(a_t.get_matrix(), w);
  return info;
};
 
 
template <typename MatrixType>
int eigh(const Tensor<MatrixType> &a, const Axes &axes_row_a,
         const Axes &axes_col_a, const Tensor<MatrixType> &b,
         const Axes &axes_row_b, const Axes &axes_col_b, std::vector<double> &w,
         Tensor<MatrixType> &z) {
  assert(axes_row_a.size() > 0);
  assert(axes_col_a.size() > 0);
  assert(axes_row_b.size() > 0);
  assert(axes_col_b.size() > 0);
 
  const size_t rank_a = a.rank();
  const size_t rank_b = b.rank();
  const size_t urank_a = axes_row_a.size();
  const size_t urank_b = axes_row_b.size();
  Axes axes_a = axes_row_a + axes_col_a;
  Axes axes_b = axes_row_b + axes_col_b;
  Tensor<MatrixType> a_t = transpose(a, axes_a, urank_a);
  Tensor<MatrixType> b_t = transpose(b, axes_b, urank_b);
  const Shape &shape_a = a_t.shape();
  const Shape &shape_b = b_t.shape();
  size_t d_row_a(1), d_col_a(1);
  size_t d_row_b(1), d_col_b(1);
  for (size_t i = 0; i < urank_a; ++i) d_row_a *= shape_a[i];
  for (size_t i = 0; i < urank_b; ++i) d_row_b *= shape_b[i];
  for (size_t i = urank_a; i < rank_a; ++i) d_col_a *= shape_a[i];
  for (size_t i = urank_b; i < rank_b; ++i) d_col_b *= shape_b[i];
 
  assert(d_row_a == d_col_a);
  assert(d_row_b == d_col_b);
  assert(d_row_a == d_row_b);
  size_t size = d_row_a;
 
  Shape shape_z;
  shape_z.resize(rank_a - urank_a + 1);
  for (size_t i = urank_a; i < rank_a; ++i) shape_z[i - urank_a] = shape_a[i];
  shape_z[rank_a - urank_a] = size;
 
  z = Tensor<MatrixType>(a.get_comm(), shape_z, urank_a);
  w.resize(size);
 
  int info;
  info = matrix_eigh(a_t.get_matrix(), b_t.get_matrix(), w, z.get_matrix());
 
  return info;
};
 
 
template <typename MatrixType>
int eig(const Tensor<MatrixType> &a, std::vector<complex> &w,
        typename Tensor<MatrixType>::template rebind<complex> &z) {
  assert(a.rank() == 2);
  Shape shape = a.shape();
  assert(shape[0] == shape[1]);
 
  Tensor<lapack::Matrix<typename MatrixType::value_type>> a_t = transpose(a, Axes(0, 1), 1).gather();
 
  size_t n = shape[0];
  w.resize(n);
 
  Tensor<lapack::Matrix<complex>> z_t(a.get_comm(), Shape(n, n), 1);
 
  int info;
  info = matrix_eig(a_t.get_matrix(), w, z_t.get_matrix());
 
  z = typename Tensor<MatrixType>::template rebind<complex>(a.get_comm(), z_t);
 
  return info;
};
 
 
template <typename MatrixType>
int eig(const Tensor<MatrixType> &a, std::vector<complex> &w) {
  assert(a.rank() == 2);
  Shape shape = a.shape();
  assert(shape[0] == shape[1]);
 
  Tensor<lapack::Matrix<typename MatrixType::value_type>> a_t = transpose(a, Axes(0, 1), 1).gather();
 
  size_t n = shape[0];
  w.resize(n);
 
  int info;
  info = matrix_eig(a_t.get_matrix(), w);
  return info;
};
 
 
template <typename MatrixType>
int eig(const Tensor<MatrixType> &a, const Axes &axes_row, const Axes &axes_col,
        std::vector<complex> &w, typename Tensor<MatrixType>::template rebind<complex> &z) {
  assert(axes_row.size() > 0);
  assert(axes_col.size() > 0);
 
  const size_t rank = a.rank();
  const size_t urank = axes_row.size();
 
  Axes axes = axes_row + axes_col;
 
  Tensor<lapack::Matrix<typename MatrixType::value_type>> a_t = transpose(a, axes, urank).gather();
  const Shape &shape = a_t.shape();
 
  size_t d_row(1), d_col(1);
  for (size_t i = 0; i < urank; ++i) d_row *= shape[i];
  for (size_t i = urank; i < rank; ++i) d_col *= shape[i];
  size_t size = (d_row < d_col) ? d_row : d_col;
 
  assert(d_row == d_col);
 
  Shape shape_z;
  shape_z.resize(urank + 1);
  for (size_t i = 0; i < urank; ++i) shape_z[i] = shape[i];
  shape_z[urank] = size;
 
  Tensor<lapack::Matrix<complex>> z_t(a.get_comm(), shape_z, urank);
 
  w.resize(size);
 
  int info;
  info = matrix_eig(a_t.get_matrix(), w, z_t.get_matrix());
 
  z = typename Tensor<MatrixType>::template rebind<complex>(a.get_comm(), z_t);
 
  return info;
};
 
 
template <typename MatrixType>
int eig(const Tensor<MatrixType> &a, const Axes &axes_row, const Axes &axes_col,
        std::vector<complex> &w) {
  assert(axes_row.size() > 0);
  assert(axes_col.size() > 0);
 
  const size_t rank = a.rank();
  const size_t urank = axes_row.size();
 
  Axes axes = axes_row + axes_col;
 
  Tensor<lapack::Matrix<typename MatrixType::value_type>> a_t = transpose(a, axes, urank).gather();
  const Shape &shape = a_t.shape();
 
  size_t d_row(1), d_col(1);
  for (size_t i = 0; i < urank; ++i) d_row *= shape[i];
  for (size_t i = urank; i < rank; ++i) d_col *= shape[i];
  size_t size = (d_row < d_col) ? d_row : d_col;
  w.resize(size);
 
  assert(d_row == d_col);
 
  int info;
  info = matrix_eig(a_t.get_matrix(), w);
  return info;
};
 
 
 
template <typename MatrixType>
int solve(const Tensor<MatrixType> &a, const std::vector<typename MatrixType::value_type> &b,
          std::vector<typename MatrixType::value_type> &x) {
  assert(a.rank() == 2);
  assert(a.shape()[0] == a.shape()[1]);
  assert(a.shape()[0] == b.size());
 
  Tensor<MatrixType> x_t;
  int info;
  info = solve(a, Tensor<MatrixType>(a.get_comm(), b), x_t, Axes(0), Axes(1),
               Axes(0), Axes());
  x = x_t.flatten();
 
  return info;
};
 
 
template <typename MatrixType>
int solve(const Tensor<MatrixType> &a, const Tensor<MatrixType> &b,
          Tensor<MatrixType> &x) {
  assert(a.rank() == 2);
  assert(b.rank() == 2 || b.rank() == 1);
  assert(a.shape()[0] == a.shape()[1]);
  assert(a.shape()[0] == b.shape()[0]);
  int info;
 
  if (b.rank() == 2) {
    info = solve(a, b, x, Axes(0), Axes(1), Axes(0), Axes(1));
  } else {
    info = solve(a, b, x, Axes(0), Axes(1), Axes(0), Axes());
  }
 
  return info;
};
 
 
template <typename MatrixType>
int solve(const Tensor<MatrixType> &a, const Tensor<MatrixType> &b,
          Tensor<MatrixType> &x, const Axes &axes_row_a, const Axes &axes_col_a,
          const Axes &axes_row_b, const Axes &axes_col_b) {
  assert(a.rank() == axes_row_a.size() + axes_col_a.size());
  assert(b.rank() == axes_row_b.size() + axes_col_b.size());
  assert(axes_row_a.size() > 0);
  assert(axes_col_a.size() > 0);
  assert(axes_row_b.size() > 0);
  assert(axes_col_b.size() >= 0);
 
  // size_t rank_a = a.rank();
  // size_t rank_b = b.rank();
  size_t rank_row_a = axes_row_a.size();
  size_t rank_col_a = axes_col_a.size();
  size_t rank_row_b = axes_row_b.size();
  size_t rank_col_b = axes_col_b.size();
  Axes axes_a = axes_row_a + axes_col_a;
  Axes axes_b = axes_row_b + axes_col_b;
 
  Tensor<MatrixType> a_t = transpose(a, axes_a, rank_row_a);
  Tensor<MatrixType> b_t = transpose(b, axes_b, rank_row_b);
  const Shape &shape_a = a_t.shape();
  const Shape &shape_b = b_t.shape();
 
  assert(debug::check_square(shape_a, rank_row_a));
 
  Shape shape_x;
  shape_x.resize(rank_col_a + rank_col_b);
  for (size_t i = 0; i < rank_col_a; ++i) shape_x[i] = shape_a[i + rank_row_a];
  for (size_t i = 0; i < rank_col_b; ++i)
    shape_x[i + rank_col_a] = shape_b[i + rank_row_b];
 
  int info;
  info = matrix_solve(a_t.get_matrix(), b_t.get_matrix());
  x = reshape(b_t, shape_x);
  return info;
};
 
template <typename MatrixType>
inline Tensor<MatrixType> operator+(Tensor<MatrixType> rhs) {
  return rhs;
}
template <typename MatrixType>
inline Tensor<MatrixType> operator-(Tensor<MatrixType> rhs) {
  return (rhs *= -1.0);
}
template <typename MatrixType>
inline Tensor<MatrixType> operator+(Tensor<MatrixType> lhs,
                                   const Tensor<MatrixType> &rhs) {
  return (lhs += rhs);
}
template <typename MatrixType>
inline Tensor<MatrixType> operator-(Tensor<MatrixType> lhs,
                                   const Tensor<MatrixType> &rhs) {
  return (lhs -= rhs);
}
template <typename MatrixType, typename D>
inline Tensor<MatrixType> operator*(Tensor<MatrixType> lhs, D rhs) {
  return (lhs *= rhs);
}
template <typename MatrixType, typename D>
inline Tensor<MatrixType> operator/(Tensor<MatrixType> lhs, D rhs) {
  return (lhs /= rhs);
}
template <typename MatrixType, typename D>
inline Tensor<MatrixType> operator*(D lhs, Tensor<MatrixType> rhs) {
  return (rhs *= lhs);
}
 
template <typename MatrixType>
inline double max(const Tensor<MatrixType> &t) {
  return max(t.get_matrix());
}
template <typename MatrixType>
inline double min(const Tensor<MatrixType> &t) {
  return min(t.get_matrix());
}
template <typename MatrixType>
inline double max_abs(const Tensor<MatrixType> &t) {
  return max_abs(t.get_matrix());
}
template <typename MatrixType>
inline double min_abs(const Tensor<MatrixType> &t) {
  return min_abs(t.get_matrix());
}
 
template <typename MatrixType>
inline Tensor<MatrixType> sqrt(Tensor<MatrixType> t) {
  using value_type = typename MatrixType::value_type;
  if constexpr (std::is_same_v<value_type, double>) {
    return t.map(static_cast<double (*)(double)>(&std::sqrt));
  } else {
    return t.map(static_cast<complex (*)(const complex &)>(&std::sqrt));
  }
}
template <typename MatrixType>
inline Tensor<MatrixType> conj(Tensor<MatrixType> t) {
  using value_type = typename MatrixType::value_type;
  if constexpr (std::is_same_v<value_type, double>) {
    return t;
  } else {
    return t.map(static_cast<complex (*)(const complex &)>(&std::conj));
  }
}
 
 
template <typename MatrixType>
std::ostream &operator<<(std::ostream &out, const Tensor<MatrixType> &t) {
  std::size_t dim = t.ndim();
  if (t.get_comm_size() == 1) {
    std::size_t size = t.local_size();
    Shape shape = t.shape();
    std::vector<std::size_t> accum(dim + 1, 1);
    for (size_t d = dim; d != 0; --d) accum[d - 1] = accum[d] * shape[d - 1];
 
    std::vector<std::size_t> idx(dim);
    for (size_t i = 0; i < size; ++i) {
      for (size_t d = 0; d < dim; ++d) idx[d] = (i % accum[d]) / accum[d + 1];
      if (idx[dim - 1] == 0) {
        size_t count = 0;
        for (size_t d = dim; d != 0 && idx[d - 1] == 0; --d) ++count;
        for (size_t i = 0; i < dim - count; ++i) out << ' ';
        for (size_t i = 0; i < count; ++i) out << '[';
      }
 
      using value_type = typename MatrixType::value_type;
      value_type val;
      t.get_value(idx, val);
      out << ' ' << val << ' ';
 
      if (idx[dim - 1] == shape[dim - 1] - 1) {
        size_t count = 0;
        for (size_t d = dim; d != 0 && idx[d - 1] == shape[d - 1] - 1; --d) ++count;
        for (size_t i = 0; i < count; ++i) out << ']';
        if (count < dim) {
          for (size_t i = 0; i < count; ++i) out << '\n';
        } else {
          out << '\n';
        }
      }
    }
  } else {
    if (t.get_comm_rank() == 0) {
      for (size_t d = 0; d < dim; ++d) out << '[';
      out << " distributed tensor ... ";
      for (size_t d = 0; d < dim; ++d) out << ']';
      out << '\n';
    }
  }
  return out;
}
 
}  // namespace mptensor
 
#endif  // _TENSOR_IMPL_HPP_