| // Copyright 2004-present Facebook. All Rights Reserved. |
| |
| #ifndef CAFFE2_OPERATORS_UTILS_EIGEN_H_ |
| #define CAFFE2_OPERATORS_UTILS_EIGEN_H_ |
| |
| #include "Eigen/Core" |
| #include "Eigen/Dense" |
| |
| #include "caffe2/core/logging.h" |
| |
| namespace caffe2 { |
| |
| // Common Eigen types that we will often use |
| template <typename T> |
| using EigenMatrixMap = |
| Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>; |
| template <typename T> |
| using EigenArrayMap = |
| Eigen::Map<Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>>; |
| template <typename T> |
| using EigenVectorMap = Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 1>>; |
| template <typename T> |
| using EigenVectorArrayMap = Eigen::Map<Eigen::Array<T, Eigen::Dynamic, 1>>; |
| template <typename T> |
| using ConstEigenMatrixMap = |
| Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>; |
| template <typename T> |
| using ConstEigenArrayMap = |
| Eigen::Map<const Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>>; |
| template <typename T> |
| using ConstEigenVectorMap = |
| Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>>; |
| template <typename T> |
| using ConstEigenVectorArrayMap = |
| Eigen::Map<const Eigen::Array<T, Eigen::Dynamic, 1>>; |
| |
| using EigenOuterStride = Eigen::OuterStride<Eigen::Dynamic>; |
| using EigenInnerStride = Eigen::InnerStride<Eigen::Dynamic>; |
| using EigenStride = Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>; |
| template <typename T> |
| using EigenOuterStridedMatrixMap = Eigen:: |
| Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>, 0, EigenOuterStride>; |
| template <typename T> |
| using EigenOuterStridedArrayMap = Eigen:: |
| Map<Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>, 0, EigenOuterStride>; |
| template <typename T> |
| using ConstEigenOuterStridedMatrixMap = Eigen::Map< |
| const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>, |
| 0, |
| EigenOuterStride>; |
| template <typename T> |
| using ConstEigenOuterStridedArrayMap = Eigen::Map< |
| const Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>, |
| 0, |
| EigenOuterStride>; |
| template <typename T> |
| using EigenStridedMatrixMap = Eigen:: |
| Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>, 0, EigenStride>; |
| template <typename T> |
| using EigenStridedArrayMap = |
| Eigen::Map<Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>, 0, EigenStride>; |
| template <typename T> |
| using ConstEigenStridedMatrixMap = Eigen:: |
| Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>, 0, EigenStride>; |
| template <typename T> |
| using ConstEigenStridedArrayMap = Eigen:: |
| Map<const Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>, 0, EigenStride>; |
| |
| // 1-d array |
| template <typename T> |
| using EArrXt = Eigen::Array<T, Eigen::Dynamic, 1>; |
| using EArrXf = Eigen::ArrayXf; |
| using EArrXd = Eigen::ArrayXd; |
| using EArrXi = Eigen::ArrayXi; |
| using EArrXb = EArrXt<bool>; |
| |
| // 2-d array, column major |
| template <typename T> |
| using EArrXXt = Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>; |
| using EArrXXf = Eigen::ArrayXXf; |
| |
| // 2-d array, row major |
| template <typename T> |
| using ERArrXXt = |
| Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>; |
| using ERArrXXf = ERArrXXt<float>; |
| |
| // 1-d vector |
| template <typename T> |
| using EVecXt = Eigen::Matrix<T, Eigen::Dynamic, 1>; |
| using EVecXd = Eigen::VectorXd; |
| using EVecXf = Eigen::VectorXf; |
| |
| // 1-d row vector |
| using ERVecXd = Eigen::RowVectorXd; |
| using ERVecXf = Eigen::RowVectorXf; |
| |
| // 2-d matrix, column major |
| template <typename T> |
| using EMatXt = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>; |
| using EMatXd = Eigen::MatrixXd; |
| using EMatXf = Eigen::MatrixXf; |
| |
| // 2-d matrix, row major |
| template <typename T> |
| using ERMatXt = |
| Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>; |
| using ERMatXd = ERMatXt<double>; |
| using ERMatXf = ERMatXt<float>; |
| |
| namespace utils { |
| |
| template <typename T> |
| Eigen::Map<const EArrXt<T>> AsEArrXt(const std::vector<T>& arr) { |
| return {arr.data(), static_cast<int>(arr.size())}; |
| } |
| template <typename T> |
| Eigen::Map<EArrXt<T>> AsEArrXt(std::vector<T>& arr) { |
| return {arr.data(), static_cast<int>(arr.size())}; |
| } |
| |
| // return a sub array of 'array' based on indices 'indices' |
| template <class Derived, class Derived1, class Derived2> |
| void GetSubArray( |
| const Eigen::ArrayBase<Derived>& array, |
| const Eigen::ArrayBase<Derived1>& indices, |
| Eigen::ArrayBase<Derived2>* out_array) { |
| CAFFE_ENFORCE_EQ(array.cols(), 1); |
| // using T = typename Derived::Scalar; |
| |
| out_array->derived().resize(indices.size()); |
| for (int i = 0; i < indices.size(); i++) { |
| DCHECK_LT(indices[i], array.size()); |
| (*out_array)[i] = array[indices[i]]; |
| } |
| } |
| |
| // return a sub array of 'array' based on indices 'indices' |
| template <class Derived, class Derived1> |
| EArrXt<typename Derived::Scalar> GetSubArray( |
| const Eigen::ArrayBase<Derived>& array, |
| const Eigen::ArrayBase<Derived1>& indices) { |
| using T = typename Derived::Scalar; |
| EArrXt<T> ret(indices.size()); |
| GetSubArray(array, indices, &ret); |
| return ret; |
| } |
| |
| // return a sub array of 'array' based on indices 'indices' |
| template <class Derived> |
| EArrXt<typename Derived::Scalar> GetSubArray( |
| const Eigen::ArrayBase<Derived>& array, |
| const std::vector<int>& indices) { |
| return GetSubArray(array, AsEArrXt(indices)); |
| } |
| |
| // return 2d sub array of 'array' based on row indices 'row_indices' |
| template <class Derived, class Derived1, class Derived2> |
| void GetSubArrayRows( |
| const Eigen::ArrayBase<Derived>& array2d, |
| const Eigen::ArrayBase<Derived1>& row_indices, |
| Eigen::ArrayBase<Derived2>* out_array) { |
| out_array->derived().resize(row_indices.size(), array2d.cols()); |
| |
| for (int i = 0; i < row_indices.size(); i++) { |
| DCHECK_LT(row_indices[i], array2d.size()); |
| out_array->row(i) = |
| array2d.row(row_indices[i]).template cast<typename Derived2::Scalar>(); |
| } |
| } |
| |
| // return indices of 1d array for elements evaluated to true |
| template <class Derived> |
| std::vector<int> GetArrayIndices(const Eigen::ArrayBase<Derived>& array) { |
| std::vector<int> ret; |
| for (int i = 0; i < array.size(); i++) { |
| if (array[i]) { |
| ret.push_back(i); |
| } |
| } |
| return ret; |
| } |
| |
| } // namespace utils |
| } // namespace caffe2 |
| |
| #endif |