matrix.h 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. /*
  3. pybind11/eigen/matrix.h: Transparent conversion for dense and sparse Eigen matrices
  4. Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
  5. All rights reserved. Use of this source code is governed by a
  6. BSD-style license that can be found in the LICENSE file.
  7. */
  8. #pragma once
  9. #include <pybind11/numpy.h>
  10. #include "common.h"
  11. /* HINT: To suppress warnings originating from the Eigen headers, use -isystem.
  12. See also:
  13. https://stackoverflow.com/questions/2579576/i-dir-vs-isystem-dir
  14. https://stackoverflow.com/questions/1741816/isystem-for-ms-visual-studio-c-compiler
  15. */
  16. PYBIND11_WARNING_PUSH
  17. PYBIND11_WARNING_DISABLE_MSVC(5054) // https://github.com/pybind/pybind11/pull/3741
  18. // C5054: operator '&': deprecated between enumerations of different types
  19. #if defined(__MINGW32__)
  20. PYBIND11_WARNING_DISABLE_GCC("-Wmaybe-uninitialized")
  21. #endif
  22. #include <Eigen/Core>
  23. #include <Eigen/SparseCore>
  24. PYBIND11_WARNING_POP
  25. // Eigen prior to 3.2.7 doesn't have proper move constructors--but worse, some classes get implicit
  26. // move constructors that break things. We could detect this an explicitly copy, but an extra copy
  27. // of matrices seems highly undesirable.
  28. static_assert(EIGEN_VERSION_AT_LEAST(3, 2, 7),
  29. "Eigen matrix support in pybind11 requires Eigen >= 3.2.7");
  30. PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
  31. PYBIND11_WARNING_DISABLE_MSVC(4127)
  32. // Provide a convenience alias for easier pass-by-ref usage with fully dynamic strides:
  33. using EigenDStride = Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>;
  34. template <typename MatrixType>
  35. using EigenDRef = Eigen::Ref<MatrixType, 0, EigenDStride>;
  36. template <typename MatrixType>
  37. using EigenDMap = Eigen::Map<MatrixType, 0, EigenDStride>;
  38. PYBIND11_NAMESPACE_BEGIN(detail)
  39. #if EIGEN_VERSION_AT_LEAST(3, 3, 0)
  40. using EigenIndex = Eigen::Index;
  41. template <typename Scalar, int Flags, typename StorageIndex>
  42. using EigenMapSparseMatrix = Eigen::Map<Eigen::SparseMatrix<Scalar, Flags, StorageIndex>>;
  43. #else
  44. using EigenIndex = EIGEN_DEFAULT_DENSE_INDEX_TYPE;
  45. template <typename Scalar, int Flags, typename StorageIndex>
  46. using EigenMapSparseMatrix = Eigen::MappedSparseMatrix<Scalar, Flags, StorageIndex>;
  47. #endif
  48. // Matches Eigen::Map, Eigen::Ref, blocks, etc:
  49. template <typename T>
  50. using is_eigen_dense_map = all_of<is_template_base_of<Eigen::DenseBase, T>,
  51. std::is_base_of<Eigen::MapBase<T, Eigen::ReadOnlyAccessors>, T>>;
  52. template <typename T>
  53. using is_eigen_mutable_map = std::is_base_of<Eigen::MapBase<T, Eigen::WriteAccessors>, T>;
  54. template <typename T>
  55. using is_eigen_dense_plain
  56. = all_of<negation<is_eigen_dense_map<T>>, is_template_base_of<Eigen::PlainObjectBase, T>>;
  57. template <typename T>
  58. using is_eigen_sparse = is_template_base_of<Eigen::SparseMatrixBase, T>;
  59. // Test for objects inheriting from EigenBase<Derived> that aren't captured by the above. This
  60. // basically covers anything that can be assigned to a dense matrix but that don't have a typical
  61. // matrix data layout that can be copied from their .data(). For example, DiagonalMatrix and
  62. // SelfAdjointView fall into this category.
  63. template <typename T>
  64. using is_eigen_other
  65. = all_of<is_template_base_of<Eigen::EigenBase, T>,
  66. negation<any_of<is_eigen_dense_map<T>, is_eigen_dense_plain<T>, is_eigen_sparse<T>>>>;
  67. // Captures numpy/eigen conformability status (returned by EigenProps::conformable()):
  68. template <bool EigenRowMajor>
  69. struct EigenConformable {
  70. bool conformable = false;
  71. EigenIndex rows = 0, cols = 0;
  72. EigenDStride stride{0, 0}; // Only valid if negativestrides is false!
  73. bool negativestrides = false; // If true, do not use stride!
  74. // NOLINTNEXTLINE(google-explicit-constructor)
  75. EigenConformable(bool fits = false) : conformable{fits} {}
  76. // Matrix type:
  77. EigenConformable(EigenIndex r, EigenIndex c, EigenIndex rstride, EigenIndex cstride)
  78. : conformable{true}, rows{r}, cols{c},
  79. // TODO: when Eigen bug #747 is fixed, remove the tests for non-negativity.
  80. // http://eigen.tuxfamily.org/bz/show_bug.cgi?id=747
  81. stride{EigenRowMajor ? (rstride > 0 ? rstride : 0)
  82. : (cstride > 0 ? cstride : 0) /* outer stride */,
  83. EigenRowMajor ? (cstride > 0 ? cstride : 0)
  84. : (rstride > 0 ? rstride : 0) /* inner stride */},
  85. negativestrides{rstride < 0 || cstride < 0} {}
  86. // Vector type:
  87. EigenConformable(EigenIndex r, EigenIndex c, EigenIndex stride)
  88. : EigenConformable(r, c, r == 1 ? c * stride : stride, c == 1 ? r : r * stride) {}
  89. template <typename props>
  90. bool stride_compatible() const {
  91. // To have compatible strides, we need (on both dimensions) one of fully dynamic strides,
  92. // matching strides, or a dimension size of 1 (in which case the stride value is
  93. // irrelevant). Alternatively, if any dimension size is 0, the strides are not relevant
  94. // (and numpy ≥ 1.23 sets the strides to 0 in that case, so we need to check explicitly).
  95. if (negativestrides) {
  96. return false;
  97. }
  98. if (rows == 0 || cols == 0) {
  99. return true;
  100. }
  101. return (props::inner_stride == Eigen::Dynamic || props::inner_stride == stride.inner()
  102. || (EigenRowMajor ? cols : rows) == 1)
  103. && (props::outer_stride == Eigen::Dynamic || props::outer_stride == stride.outer()
  104. || (EigenRowMajor ? rows : cols) == 1);
  105. }
  106. // NOLINTNEXTLINE(google-explicit-constructor)
  107. operator bool() const { return conformable; }
  108. };
  109. template <typename Type>
  110. struct eigen_extract_stride {
  111. using type = Type;
  112. };
  113. template <typename PlainObjectType, int MapOptions, typename StrideType>
  114. struct eigen_extract_stride<Eigen::Map<PlainObjectType, MapOptions, StrideType>> {
  115. using type = StrideType;
  116. };
  117. template <typename PlainObjectType, int Options, typename StrideType>
  118. struct eigen_extract_stride<Eigen::Ref<PlainObjectType, Options, StrideType>> {
  119. using type = StrideType;
  120. };
  121. // Helper struct for extracting information from an Eigen type
  122. template <typename Type_>
  123. struct EigenProps {
  124. using Type = Type_;
  125. using Scalar = typename Type::Scalar;
  126. using StrideType = typename eigen_extract_stride<Type>::type;
  127. static constexpr EigenIndex rows = Type::RowsAtCompileTime, cols = Type::ColsAtCompileTime,
  128. size = Type::SizeAtCompileTime;
  129. static constexpr bool row_major = Type::IsRowMajor,
  130. vector
  131. = Type::IsVectorAtCompileTime, // At least one dimension has fixed size 1
  132. fixed_rows = rows != Eigen::Dynamic, fixed_cols = cols != Eigen::Dynamic,
  133. fixed = size != Eigen::Dynamic, // Fully-fixed size
  134. dynamic = !fixed_rows && !fixed_cols; // Fully-dynamic size
  135. template <EigenIndex i, EigenIndex ifzero>
  136. using if_zero = std::integral_constant<EigenIndex, i == 0 ? ifzero : i>;
  137. static constexpr EigenIndex inner_stride
  138. = if_zero<StrideType::InnerStrideAtCompileTime, 1>::value,
  139. outer_stride = if_zero < StrideType::OuterStrideAtCompileTime,
  140. vector ? size
  141. : row_major ? cols
  142. : rows > ::value;
  143. static constexpr bool dynamic_stride
  144. = inner_stride == Eigen::Dynamic && outer_stride == Eigen::Dynamic;
  145. static constexpr bool requires_row_major
  146. = !dynamic_stride && !vector && (row_major ? inner_stride : outer_stride) == 1;
  147. static constexpr bool requires_col_major
  148. = !dynamic_stride && !vector && (row_major ? outer_stride : inner_stride) == 1;
  149. // Takes an input array and determines whether we can make it fit into the Eigen type. If
  150. // the array is a vector, we attempt to fit it into either an Eigen 1xN or Nx1 vector
  151. // (preferring the latter if it will fit in either, i.e. for a fully dynamic matrix type).
  152. static EigenConformable<row_major> conformable(const array &a) {
  153. const auto dims = a.ndim();
  154. if (dims < 1 || dims > 2) {
  155. return false;
  156. }
  157. if (dims == 2) { // Matrix type: require exact match (or dynamic)
  158. EigenIndex np_rows = a.shape(0), np_cols = a.shape(1),
  159. np_rstride = a.strides(0) / static_cast<ssize_t>(sizeof(Scalar)),
  160. np_cstride = a.strides(1) / static_cast<ssize_t>(sizeof(Scalar));
  161. if ((fixed_rows && np_rows != rows) || (fixed_cols && np_cols != cols)) {
  162. return false;
  163. }
  164. return {np_rows, np_cols, np_rstride, np_cstride};
  165. }
  166. // Otherwise we're storing an n-vector. Only one of the strides will be used, but
  167. // whichever is used, we want the (single) numpy stride value.
  168. const EigenIndex n = a.shape(0),
  169. stride = a.strides(0) / static_cast<ssize_t>(sizeof(Scalar));
  170. if (vector) { // Eigen type is a compile-time vector
  171. if (fixed && size != n) {
  172. return false; // Vector size mismatch
  173. }
  174. return {rows == 1 ? 1 : n, cols == 1 ? 1 : n, stride};
  175. }
  176. if (fixed) {
  177. // The type has a fixed size, but is not a vector: abort
  178. return false;
  179. }
  180. if (fixed_cols) {
  181. // Since this isn't a vector, cols must be != 1. We allow this only if it exactly
  182. // equals the number of elements (rows is Dynamic, and so 1 row is allowed).
  183. if (cols != n) {
  184. return false;
  185. }
  186. return {1, n, stride};
  187. } // Otherwise it's either fully dynamic, or column dynamic; both become a column vector
  188. if (fixed_rows && rows != n) {
  189. return false;
  190. }
  191. return {n, 1, stride};
  192. }
  193. static constexpr bool show_writeable
  194. = is_eigen_dense_map<Type>::value && is_eigen_mutable_map<Type>::value;
  195. static constexpr bool show_order = is_eigen_dense_map<Type>::value;
  196. static constexpr bool show_c_contiguous = show_order && requires_row_major;
  197. static constexpr bool show_f_contiguous
  198. = !show_c_contiguous && show_order && requires_col_major;
  199. static constexpr auto descriptor
  200. = const_name("typing.Annotated[")
  201. + io_name("numpy.typing.ArrayLike, ", "numpy.typing.NDArray[")
  202. + npy_format_descriptor<Scalar>::name + io_name("", "]") + const_name(", \"[")
  203. + const_name<fixed_rows>(const_name<(size_t) rows>(), const_name("m")) + const_name(", ")
  204. + const_name<fixed_cols>(const_name<(size_t) cols>(), const_name("n"))
  205. + const_name("]\"")
  206. // For a reference type (e.g. Ref<MatrixXd>) we have other constraints that might need to
  207. // be satisfied: writeable=True (for a mutable reference), and, depending on the map's
  208. // stride options, possibly f_contiguous or c_contiguous. We include them in the
  209. // descriptor output to provide some hint as to why a TypeError is occurring (otherwise
  210. // it can be confusing to see that a function accepts a
  211. // 'typing.Annotated[numpy.typing.NDArray[numpy.float64], "[3,2]"]' and an error message
  212. // that you *gave* a numpy.ndarray of the right type and dimensions.
  213. + const_name<show_writeable>(", \"flags.writeable\"", "")
  214. + const_name<show_c_contiguous>(", \"flags.c_contiguous\"", "")
  215. + const_name<show_f_contiguous>(", \"flags.f_contiguous\"", "") + const_name("]");
  216. };
  217. // Casts an Eigen type to numpy array. If given a base, the numpy array references the src data,
  218. // otherwise it'll make a copy. writeable lets you turn off the writeable flag for the array.
  219. template <typename props>
  220. handle
  221. eigen_array_cast(typename props::Type const &src, handle base = handle(), bool writeable = true) {
  222. constexpr ssize_t elem_size = sizeof(typename props::Scalar);
  223. array a;
  224. if (props::vector) {
  225. a = array({src.size()}, {elem_size * src.innerStride()}, src.data(), base);
  226. } else {
  227. a = array({src.rows(), src.cols()},
  228. {elem_size * src.rowStride(), elem_size * src.colStride()},
  229. src.data(),
  230. base);
  231. }
  232. if (!writeable) {
  233. array_proxy(a.ptr())->flags &= ~detail::npy_api::NPY_ARRAY_WRITEABLE_;
  234. }
  235. return a.release();
  236. }
  237. // Takes an lvalue ref to some Eigen type and a (python) base object, creating a numpy array that
  238. // reference the Eigen object's data with `base` as the python-registered base class (if omitted,
  239. // the base will be set to None, and lifetime management is up to the caller). The numpy array is
  240. // non-writeable if the given type is const.
  241. template <typename props, typename Type>
  242. handle eigen_ref_array(Type &src, handle parent = none()) {
  243. // none here is to get past array's should-we-copy detection, which currently always
  244. // copies when there is no base. Setting the base to None should be harmless.
  245. return eigen_array_cast<props>(src, parent, !std::is_const<Type>::value);
  246. }
  247. // Takes a pointer to some dense, plain Eigen type, builds a capsule around it, then returns a
  248. // numpy array that references the encapsulated data with a python-side reference to the capsule to
  249. // tie its destruction to that of any dependent python objects. Const-ness is determined by
  250. // whether or not the Type of the pointer given is const.
  251. template <typename props, typename Type, typename = enable_if_t<is_eigen_dense_plain<Type>::value>>
  252. handle eigen_encapsulate(Type *src) {
  253. capsule base(src, [](void *o) { delete static_cast<Type *>(o); });
  254. return eigen_ref_array<props>(*src, base);
  255. }
  256. // Type caster for regular, dense matrix types (e.g. MatrixXd), but not maps/refs/etc. of dense
  257. // types.
  258. template <typename Type>
  259. struct type_caster<Type, enable_if_t<is_eigen_dense_plain<Type>::value>> {
  260. using Scalar = typename Type::Scalar;
  261. static_assert(!std::is_pointer<Scalar>::value,
  262. PYBIND11_EIGEN_MESSAGE_POINTER_TYPES_ARE_NOT_SUPPORTED);
  263. using props = EigenProps<Type>;
  264. bool load(handle src, bool convert) {
  265. // If we're in no-convert mode, only load if given an array of the correct type
  266. if (!convert && !isinstance<array_t<Scalar>>(src)) {
  267. return false;
  268. }
  269. // Coerce into an array, but don't do type conversion yet; the copy below handles it.
  270. auto buf = array::ensure(src);
  271. if (!buf) {
  272. return false;
  273. }
  274. auto dims = buf.ndim();
  275. if (dims < 1 || dims > 2) {
  276. return false;
  277. }
  278. auto fits = props::conformable(buf);
  279. if (!fits) {
  280. return false;
  281. }
  282. PYBIND11_WARNING_PUSH
  283. PYBIND11_WARNING_DISABLE_GCC("-Wmaybe-uninitialized") // See PR #5516
  284. // Allocate the new type, then build a numpy reference into it
  285. value = Type(fits.rows, fits.cols);
  286. PYBIND11_WARNING_POP
  287. auto ref = reinterpret_steal<array>(eigen_ref_array<props>(value));
  288. if (dims == 1) {
  289. ref = ref.squeeze();
  290. } else if (ref.ndim() == 1) {
  291. buf = buf.squeeze();
  292. }
  293. int result = detail::npy_api::get().PyArray_CopyInto_(ref.ptr(), buf.ptr());
  294. if (result < 0) { // Copy failed!
  295. PyErr_Clear();
  296. return false;
  297. }
  298. return true;
  299. }
  300. private:
  301. // Cast implementation
  302. template <typename CType>
  303. static handle cast_impl(CType *src, return_value_policy policy, handle parent) {
  304. switch (policy) {
  305. case return_value_policy::take_ownership:
  306. case return_value_policy::automatic:
  307. return eigen_encapsulate<props>(src);
  308. case return_value_policy::move:
  309. return eigen_encapsulate<props>(new CType(std::move(*src)));
  310. case return_value_policy::copy:
  311. return eigen_array_cast<props>(*src);
  312. case return_value_policy::reference:
  313. case return_value_policy::automatic_reference:
  314. return eigen_ref_array<props>(*src);
  315. case return_value_policy::reference_internal:
  316. return eigen_ref_array<props>(*src, parent);
  317. default:
  318. throw cast_error("unhandled return_value_policy: should not happen!");
  319. };
  320. }
  321. public:
  322. // Normal returned non-reference, non-const value:
  323. static handle cast(Type &&src, return_value_policy /* policy */, handle parent) {
  324. return cast_impl(&src, return_value_policy::move, parent);
  325. }
  326. // If you return a non-reference const, we mark the numpy array readonly:
  327. static handle cast(const Type &&src, return_value_policy /* policy */, handle parent) {
  328. return cast_impl(&src, return_value_policy::move, parent);
  329. }
  330. // lvalue reference return; default (automatic) becomes copy
  331. static handle cast(Type &src, return_value_policy policy, handle parent) {
  332. if (policy == return_value_policy::automatic
  333. || policy == return_value_policy::automatic_reference) {
  334. policy = return_value_policy::copy;
  335. }
  336. return cast_impl(&src, policy, parent);
  337. }
  338. // const lvalue reference return; default (automatic) becomes copy
  339. static handle cast(const Type &src, return_value_policy policy, handle parent) {
  340. if (policy == return_value_policy::automatic
  341. || policy == return_value_policy::automatic_reference) {
  342. policy = return_value_policy::copy;
  343. }
  344. return cast(&src, policy, parent);
  345. }
  346. // non-const pointer return
  347. static handle cast(Type *src, return_value_policy policy, handle parent) {
  348. return cast_impl(src, policy, parent);
  349. }
  350. // const pointer return
  351. static handle cast(const Type *src, return_value_policy policy, handle parent) {
  352. return cast_impl(src, policy, parent);
  353. }
  354. static constexpr auto name = props::descriptor;
  355. // NOLINTNEXTLINE(google-explicit-constructor)
  356. operator Type *() { return &value; }
  357. // NOLINTNEXTLINE(google-explicit-constructor)
  358. operator Type &() { return value; }
  359. // NOLINTNEXTLINE(google-explicit-constructor)
  360. operator Type &&() && { return std::move(value); }
  361. template <typename T>
  362. using cast_op_type = movable_cast_op_type<T>;
  363. private:
  364. Type value;
  365. };
  366. // Base class for casting reference/map/block/etc. objects back to python.
  367. template <typename MapType>
  368. struct eigen_map_caster {
  369. static_assert(!std::is_pointer<typename MapType::Scalar>::value,
  370. PYBIND11_EIGEN_MESSAGE_POINTER_TYPES_ARE_NOT_SUPPORTED);
  371. private:
  372. using props = EigenProps<MapType>;
  373. public:
  374. // Directly referencing a ref/map's data is a bit dangerous (whatever the map/ref points to has
  375. // to stay around), but we'll allow it under the assumption that you know what you're doing
  376. // (and have an appropriate keep_alive in place). We return a numpy array pointing directly at
  377. // the ref's data (The numpy array ends up read-only if the ref was to a const matrix type.)
  378. // Note that this means you need to ensure you don't destroy the object in some other way (e.g.
  379. // with an appropriate keep_alive, or with a reference to a statically allocated matrix).
  380. static handle cast(const MapType &src, return_value_policy policy, handle parent) {
  381. switch (policy) {
  382. case return_value_policy::copy:
  383. return eigen_array_cast<props>(src);
  384. case return_value_policy::reference_internal:
  385. return eigen_array_cast<props>(src, parent, is_eigen_mutable_map<MapType>::value);
  386. case return_value_policy::reference:
  387. case return_value_policy::automatic:
  388. case return_value_policy::automatic_reference:
  389. return eigen_array_cast<props>(src, none(), is_eigen_mutable_map<MapType>::value);
  390. default:
  391. // move, take_ownership don't make any sense for a ref/map:
  392. pybind11_fail("Invalid return_value_policy for Eigen Map/Ref/Block type");
  393. }
  394. }
  395. // return_descr forces the use of NDArray instead of ArrayLike in args
  396. // since Ref<...> args can only accept arrays.
  397. static constexpr auto name = return_descr(props::descriptor);
  398. // Explicitly delete these: support python -> C++ conversion on these (i.e. these can be return
  399. // types but not bound arguments). We still provide them (with an explicitly delete) so that
  400. // you end up here if you try anyway.
  401. bool load(handle, bool) = delete;
  402. operator MapType() = delete;
  403. template <typename>
  404. using cast_op_type = MapType;
  405. };
  406. // We can return any map-like object (but can only load Refs, specialized next):
  407. template <typename Type>
  408. struct type_caster<Type, enable_if_t<is_eigen_dense_map<Type>::value>> : eigen_map_caster<Type> {};
  409. // Loader for Ref<...> arguments. See the documentation for info on how to make this work without
  410. // copying (it requires some extra effort in many cases).
  411. template <typename PlainObjectType, typename StrideType>
  412. struct type_caster<
  413. Eigen::Ref<PlainObjectType, 0, StrideType>,
  414. enable_if_t<is_eigen_dense_map<Eigen::Ref<PlainObjectType, 0, StrideType>>::value>>
  415. : public eigen_map_caster<Eigen::Ref<PlainObjectType, 0, StrideType>> {
  416. private:
  417. using Type = Eigen::Ref<PlainObjectType, 0, StrideType>;
  418. using props = EigenProps<Type>;
  419. using Scalar = typename props::Scalar;
  420. static_assert(!std::is_pointer<Scalar>::value,
  421. PYBIND11_EIGEN_MESSAGE_POINTER_TYPES_ARE_NOT_SUPPORTED);
  422. using MapType = Eigen::Map<PlainObjectType, 0, StrideType>;
  423. using Array
  424. = array_t<Scalar,
  425. array::forcecast
  426. | ((props::row_major ? props::inner_stride : props::outer_stride) == 1
  427. ? array::c_style
  428. : (props::row_major ? props::outer_stride : props::inner_stride) == 1
  429. ? array::f_style
  430. : 0)>;
  431. static constexpr bool need_writeable = is_eigen_mutable_map<Type>::value;
  432. // Delay construction (these have no default constructor)
  433. std::unique_ptr<MapType> map;
  434. std::unique_ptr<Type> ref;
  435. // Our array. When possible, this is just a numpy array pointing to the source data, but
  436. // sometimes we can't avoid copying (e.g. input is not a numpy array at all, has an
  437. // incompatible layout, or is an array of a type that needs to be converted). Using a numpy
  438. // temporary (rather than an Eigen temporary) saves an extra copy when we need both type
  439. // conversion and storage order conversion. (Note that we refuse to use this temporary copy
  440. // when loading an argument for a Ref<M> with M non-const, i.e. a read-write reference).
  441. Array copy_or_ref;
  442. public:
  443. bool load(handle src, bool convert) {
  444. // First check whether what we have is already an array of the right type. If not, we
  445. // can't avoid a copy (because the copy is also going to do type conversion).
  446. bool need_copy = !isinstance<Array>(src);
  447. EigenConformable<props::row_major> fits;
  448. if (!need_copy) {
  449. // We don't need a converting copy, but we also need to check whether the strides are
  450. // compatible with the Ref's stride requirements
  451. auto aref = reinterpret_borrow<Array>(src);
  452. if (aref && (!need_writeable || aref.writeable())) {
  453. fits = props::conformable(aref);
  454. if (!fits) {
  455. return false; // Incompatible dimensions
  456. }
  457. if (!fits.template stride_compatible<props>()) {
  458. need_copy = true;
  459. } else {
  460. copy_or_ref = std::move(aref);
  461. }
  462. } else {
  463. need_copy = true;
  464. }
  465. }
  466. if (need_copy) {
  467. // We need to copy: If we need a mutable reference, or we're not supposed to convert
  468. // (either because we're in the no-convert overload pass, or because we're explicitly
  469. // instructed not to copy (via `py::arg().noconvert()`) we have to fail loading.
  470. if (!convert || need_writeable) {
  471. return false;
  472. }
  473. Array copy = Array::ensure(src);
  474. if (!copy) {
  475. return false;
  476. }
  477. fits = props::conformable(copy);
  478. if (!fits || !fits.template stride_compatible<props>()) {
  479. return false;
  480. }
  481. copy_or_ref = std::move(copy);
  482. loader_life_support::add_patient(copy_or_ref);
  483. }
  484. ref.reset();
  485. map.reset(new MapType(data(copy_or_ref),
  486. fits.rows,
  487. fits.cols,
  488. make_stride(fits.stride.outer(), fits.stride.inner())));
  489. ref.reset(new Type(*map));
  490. return true;
  491. }
  492. // NOLINTNEXTLINE(google-explicit-constructor)
  493. operator Type *() { return ref.get(); }
  494. // NOLINTNEXTLINE(google-explicit-constructor)
  495. operator Type &() { return *ref; }
  496. template <typename _T>
  497. using cast_op_type = pybind11::detail::cast_op_type<_T>;
  498. private:
  499. template <typename T = Type, enable_if_t<is_eigen_mutable_map<T>::value, int> = 0>
  500. Scalar *data(Array &a) {
  501. return a.mutable_data();
  502. }
  503. template <typename T = Type, enable_if_t<!is_eigen_mutable_map<T>::value, int> = 0>
  504. const Scalar *data(Array &a) {
  505. return a.data();
  506. }
  507. // Attempt to figure out a constructor of `Stride` that will work.
  508. // If both strides are fixed, use a default constructor:
  509. template <typename S>
  510. using stride_ctor_default = bool_constant<S::InnerStrideAtCompileTime != Eigen::Dynamic
  511. && S::OuterStrideAtCompileTime != Eigen::Dynamic
  512. && std::is_default_constructible<S>::value>;
  513. // Otherwise, if there is a two-index constructor, assume it is (outer,inner) like
  514. // Eigen::Stride, and use it:
  515. template <typename S>
  516. using stride_ctor_dual
  517. = bool_constant<!stride_ctor_default<S>::value
  518. && std::is_constructible<S, EigenIndex, EigenIndex>::value>;
  519. // Otherwise, if there is a one-index constructor, and just one of the strides is dynamic, use
  520. // it (passing whichever stride is dynamic).
  521. template <typename S>
  522. using stride_ctor_outer
  523. = bool_constant<!any_of<stride_ctor_default<S>, stride_ctor_dual<S>>::value
  524. && S::OuterStrideAtCompileTime == Eigen::Dynamic
  525. && S::InnerStrideAtCompileTime != Eigen::Dynamic
  526. && std::is_constructible<S, EigenIndex>::value>;
  527. template <typename S>
  528. using stride_ctor_inner
  529. = bool_constant<!any_of<stride_ctor_default<S>, stride_ctor_dual<S>>::value
  530. && S::InnerStrideAtCompileTime == Eigen::Dynamic
  531. && S::OuterStrideAtCompileTime != Eigen::Dynamic
  532. && std::is_constructible<S, EigenIndex>::value>;
  533. template <typename S = StrideType, enable_if_t<stride_ctor_default<S>::value, int> = 0>
  534. static S make_stride(EigenIndex, EigenIndex) {
  535. return S();
  536. }
  537. template <typename S = StrideType, enable_if_t<stride_ctor_dual<S>::value, int> = 0>
  538. static S make_stride(EigenIndex outer, EigenIndex inner) {
  539. return S(outer, inner);
  540. }
  541. template <typename S = StrideType, enable_if_t<stride_ctor_outer<S>::value, int> = 0>
  542. static S make_stride(EigenIndex outer, EigenIndex) {
  543. return S(outer);
  544. }
  545. template <typename S = StrideType, enable_if_t<stride_ctor_inner<S>::value, int> = 0>
  546. static S make_stride(EigenIndex, EigenIndex inner) {
  547. return S(inner);
  548. }
  549. };
  550. // type_caster for special matrix types (e.g. DiagonalMatrix), which are EigenBase, but not
  551. // EigenDense (i.e. they don't have a data(), at least not with the usual matrix layout).
  552. // load() is not supported, but we can cast them into the python domain by first copying to a
  553. // regular Eigen::Matrix, then casting that.
  554. template <typename Type>
  555. struct type_caster<Type, enable_if_t<is_eigen_other<Type>::value>> {
  556. static_assert(!std::is_pointer<typename Type::Scalar>::value,
  557. PYBIND11_EIGEN_MESSAGE_POINTER_TYPES_ARE_NOT_SUPPORTED);
  558. protected:
  559. using Matrix
  560. = Eigen::Matrix<typename Type::Scalar, Type::RowsAtCompileTime, Type::ColsAtCompileTime>;
  561. using props = EigenProps<Matrix>;
  562. public:
  563. static handle cast(const Type &src, return_value_policy /* policy */, handle /* parent */) {
  564. handle h = eigen_encapsulate<props>(new Matrix(src));
  565. return h;
  566. }
  567. static handle cast(const Type *src, return_value_policy policy, handle parent) {
  568. return cast(*src, policy, parent);
  569. }
  570. static constexpr auto name = props::descriptor;
  571. // Explicitly delete these: support python -> C++ conversion on these (i.e. these can be return
  572. // types but not bound arguments). We still provide them (with an explicitly delete) so that
  573. // you end up here if you try anyway.
  574. bool load(handle, bool) = delete;
  575. operator Type() = delete;
  576. template <typename>
  577. using cast_op_type = Type;
  578. };
  579. template <typename Type>
  580. struct type_caster<Type, enable_if_t<is_eigen_sparse<Type>::value>> {
  581. using Scalar = typename Type::Scalar;
  582. static_assert(!std::is_pointer<Scalar>::value,
  583. PYBIND11_EIGEN_MESSAGE_POINTER_TYPES_ARE_NOT_SUPPORTED);
  584. using StorageIndex = remove_reference_t<decltype(*std::declval<Type>().outerIndexPtr())>;
  585. using Index = typename Type::Index;
  586. static constexpr bool rowMajor = Type::IsRowMajor;
  587. bool load(handle src, bool) {
  588. if (!src) {
  589. return false;
  590. }
  591. auto obj = reinterpret_borrow<object>(src);
  592. object sparse_module = module_::import("scipy.sparse");
  593. object matrix_type = sparse_module.attr(rowMajor ? "csr_matrix" : "csc_matrix");
  594. if (!type::handle_of(obj).is(matrix_type)) {
  595. try {
  596. obj = matrix_type(obj);
  597. } catch (const error_already_set &) {
  598. return false;
  599. }
  600. }
  601. auto values = array_t<Scalar>((object) obj.attr("data"));
  602. auto innerIndices = array_t<StorageIndex>((object) obj.attr("indices"));
  603. auto outerIndices = array_t<StorageIndex>((object) obj.attr("indptr"));
  604. auto shape = pybind11::tuple((pybind11::object) obj.attr("shape"));
  605. auto nnz = obj.attr("nnz").cast<Index>();
  606. if (!values || !innerIndices || !outerIndices) {
  607. return false;
  608. }
  609. value = EigenMapSparseMatrix<Scalar,
  610. Type::Flags &(Eigen::RowMajor | Eigen::ColMajor),
  611. StorageIndex>(shape[0].cast<Index>(),
  612. shape[1].cast<Index>(),
  613. std::move(nnz),
  614. outerIndices.mutable_data(),
  615. innerIndices.mutable_data(),
  616. values.mutable_data());
  617. return true;
  618. }
  619. static handle cast(const Type &src, return_value_policy /* policy */, handle /* parent */) {
  620. const_cast<Type &>(src).makeCompressed();
  621. object matrix_type
  622. = module_::import("scipy.sparse").attr(rowMajor ? "csr_matrix" : "csc_matrix");
  623. array data(src.nonZeros(), src.valuePtr());
  624. array outerIndices((rowMajor ? src.rows() : src.cols()) + 1, src.outerIndexPtr());
  625. array innerIndices(src.nonZeros(), src.innerIndexPtr());
  626. return matrix_type(pybind11::make_tuple(
  627. std::move(data), std::move(innerIndices), std::move(outerIndices)),
  628. pybind11::make_tuple(src.rows(), src.cols()))
  629. .release();
  630. }
  631. PYBIND11_TYPE_CASTER(Type,
  632. const_name<(Type::IsRowMajor) != 0>("scipy.sparse.csr_matrix[",
  633. "scipy.sparse.csc_matrix[")
  634. + npy_format_descriptor<Scalar>::name + const_name("]"));
  635. };
  636. PYBIND11_NAMESPACE_END(detail)
  637. PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)
  638. #else
  639. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  640. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)