pybind11/eigen/tensor.h

0001 /*
0002     pybind11/eigen/tensor.h: Transparent conversion for Eigen tensors
0003
0004     All rights reserved. Use of this source code is governed by a
0005     BSD-style license that can be found in the LICENSE file.
0006 */
0007
0008 #pragma once
0009
0010 #include <pybind11/numpy.h>
0011
0012 #include "common.h"
0013
0014 #if defined(__GNUC__) && !defined(__clang__) && !defined(__INTEL_COMPILER)
0015 static_assert(__GNUC__ > 5, "Eigen Tensor support in pybind11 requires GCC > 5.0");
0016 #endif
0017
0018 // Disable warnings for Eigen
0019 PYBIND11_WARNING_PUSH
0020 PYBIND11_WARNING_DISABLE_MSVC(4554)
0021 PYBIND11_WARNING_DISABLE_MSVC(4127)
0022 #if defined(__MINGW32__)
0023 PYBIND11_WARNING_DISABLE_GCC("-Wmaybe-uninitialized")
0024 #endif
0025
0026 #include <unsupported/Eigen/CXX11/Tensor>
0027
0028 PYBIND11_WARNING_POP
0029
0030 static_assert(EIGEN_VERSION_AT_LEAST(3, 3, 0),
0031               "Eigen Tensor support in pybind11 requires Eigen >= 3.3.0");
0032
0033 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
0034
0035 PYBIND11_WARNING_DISABLE_MSVC(4127)
0036
0037 PYBIND11_NAMESPACE_BEGIN(detail)
0038
0039 inline bool is_tensor_aligned(const void *data) {
0040     return (reinterpret_cast<std::size_t>(data) % EIGEN_DEFAULT_ALIGN_BYTES) == 0;
0041 }
0042
0043 template <typename T>
0044 constexpr int compute_array_flag_from_tensor() {
0045     static_assert((static_cast<int>(T::Layout) == static_cast<int>(Eigen::RowMajor))
0046                       || (static_cast<int>(T::Layout) == static_cast<int>(Eigen::ColMajor)),
0047                   "Layout must be row or column major");
0048     return (static_cast<int>(T::Layout) == static_cast<int>(Eigen::RowMajor)) ? array::c_style
0049                                                                               : array::f_style;
0050 }
0051
0052 template <typename T>
0053 struct eigen_tensor_helper {};
0054
0055 template <typename Scalar_, int NumIndices_, int Options_, typename IndexType>
0056 struct eigen_tensor_helper<Eigen::Tensor<Scalar_, NumIndices_, Options_, IndexType>> {
0057     using Type = Eigen::Tensor<Scalar_, NumIndices_, Options_, IndexType>;
0058     using ValidType = void;
0059
0060     static Eigen::DSizes<typename Type::Index, Type::NumIndices> get_shape(const Type &f) {
0061         return f.dimensions();
0062     }
0063
0064     static constexpr bool
0065     is_correct_shape(const Eigen::DSizes<typename Type::Index, Type::NumIndices> & /*shape*/) {
0066         return true;
0067     }
0068
0069     template <typename T>
0070     struct helper {};
0071
0072     template <size_t... Is>
0073     struct helper<index_sequence<Is...>> {
0074         static constexpr auto value = ::pybind11::detail::concat(const_name(((void) Is, "?"))...);
0075     };
0076
0077     static constexpr auto dimensions_descriptor
0078         = helper<decltype(make_index_sequence<Type::NumIndices>())>::value;
0079
0080     template <typename... Args>
0081     static Type *alloc(Args &&...args) {
0082         return new Type(std::forward<Args>(args)...);
0083     }
0084
0085     static void free(Type *tensor) { delete tensor; }
0086 };
0087
0088 template <typename Scalar_, typename std::ptrdiff_t... Indices, int Options_, typename IndexType>
0089 struct eigen_tensor_helper<
0090     Eigen::TensorFixedSize<Scalar_, Eigen::Sizes<Indices...>, Options_, IndexType>> {
0091     using Type = Eigen::TensorFixedSize<Scalar_, Eigen::Sizes<Indices...>, Options_, IndexType>;
0092     using ValidType = void;
0093
0094     static constexpr Eigen::DSizes<typename Type::Index, Type::NumIndices>
0095     get_shape(const Type & /*f*/) {
0096         return get_shape();
0097     }
0098
0099     static constexpr Eigen::DSizes<typename Type::Index, Type::NumIndices> get_shape() {
0100         return Eigen::DSizes<typename Type::Index, Type::NumIndices>(Indices...);
0101     }
0102
0103     static bool
0104     is_correct_shape(const Eigen::DSizes<typename Type::Index, Type::NumIndices> &shape) {
0105         return get_shape() == shape;
0106     }
0107
0108     static constexpr auto dimensions_descriptor
0109         = ::pybind11::detail::concat(const_name<Indices>()...);
0110
0111     template <typename... Args>
0112     static Type *alloc(Args &&...args) {
0113         Eigen::aligned_allocator<Type> allocator;
0114         return ::new (allocator.allocate(1)) Type(std::forward<Args>(args)...);
0115     }
0116
0117     static void free(Type *tensor) {
0118         Eigen::aligned_allocator<Type> allocator;
0119         tensor->~Type();
0120         allocator.deallocate(tensor, 1);
0121     }
0122 };
0123
0124 template <typename Type, bool ShowDetails, bool NeedsWriteable = false>
0125 struct get_tensor_descriptor {
0126     static constexpr auto details
0127         = const_name<NeedsWriteable>(", \"flags.writeable\"", "")
0128           + const_name<static_cast<int>(Type::Layout) == static_cast<int>(Eigen::RowMajor)>(
0129               ", \"flags.c_contiguous\"", ", \"flags.f_contiguous\"");
0130     static constexpr auto value
0131         = const_name("typing.Annotated[")
0132           + io_name("numpy.typing.ArrayLike, ", "numpy.typing.NDArray[")
0133           + npy_format_descriptor<typename Type::Scalar>::name + io_name("", "]")
0134           + const_name(", \"[") + eigen_tensor_helper<remove_cv_t<Type>>::dimensions_descriptor
0135           + const_name("]\"") + const_name<ShowDetails>(details, const_name("")) + const_name("]");
0136 };
0137
0138 // When EIGEN_AVOID_STL_ARRAY is defined, Eigen::DSizes<T, 0> does not have the begin() member
0139 // function. Falling back to a simple loop works around this issue.
0140 //
0141 // We need to disable the type-limits warning for the inner loop when size = 0.
0142
0143 PYBIND11_WARNING_PUSH
0144 PYBIND11_WARNING_DISABLE_GCC("-Wtype-limits")
0145
0146 template <typename T, int size>
0147 std::vector<T> convert_dsizes_to_vector(const Eigen::DSizes<T, size> &arr) {
0148     std::vector<T> result(size);
0149
0150     for (size_t i = 0; i < size; i++) {
0151         result[i] = arr[i];
0152     }
0153
0154     return result;
0155 }
0156
0157 template <typename T, int size>
0158 Eigen::DSizes<T, size> get_shape_for_array(const array &arr) {
0159     Eigen::DSizes<T, size> result;
0160     const T *shape = arr.shape();
0161     for (size_t i = 0; i < size; i++) {
0162         result[i] = shape[i];
0163     }
0164
0165     return result;
0166 }
0167
0168 PYBIND11_WARNING_POP
0169
0170 template <typename Type>
0171 struct type_caster<Type, typename eigen_tensor_helper<Type>::ValidType> {
0172     static_assert(!std::is_pointer<typename Type::Scalar>::value,
0173                   PYBIND11_EIGEN_MESSAGE_POINTER_TYPES_ARE_NOT_SUPPORTED);
0174     using Helper = eigen_tensor_helper<Type>;
0175     static constexpr auto temp_name = get_tensor_descriptor<Type, false>::value;
0176     PYBIND11_TYPE_CASTER(Type, temp_name);
0177
0178     bool load(handle src, bool convert) {
0179         if (!convert) {
0180             if (!isinstance<array>(src)) {
0181                 return false;
0182             }
0183             array temp = array::ensure(src);
0184             if (!temp) {
0185                 return false;
0186             }
0187
0188             if (!temp.dtype().is(dtype::of<typename Type::Scalar>())) {
0189                 return false;
0190             }
0191         }
0192
0193         array_t<typename Type::Scalar, compute_array_flag_from_tensor<Type>()> arr(
0194             reinterpret_borrow<object>(src));
0195
0196         if (arr.ndim() != Type::NumIndices) {
0197             return false;
0198         }
0199         auto shape = get_shape_for_array<typename Type::Index, Type::NumIndices>(arr);
0200
0201         if (!Helper::is_correct_shape(shape)) {
0202             return false;
0203         }
0204
0205 #if EIGEN_VERSION_AT_LEAST(3, 4, 0)
0206         auto data_pointer = arr.data();
0207 #else
0208         // Handle Eigen bug
0209         auto data_pointer = const_cast<typename Type::Scalar *>(arr.data());
0210 #endif
0211
0212         if (is_tensor_aligned(arr.data())) {
0213             value = Eigen::TensorMap<const Type, Eigen::Aligned>(data_pointer, shape);
0214         } else {
0215             value = Eigen::TensorMap<const Type>(data_pointer, shape);
0216         }
0217
0218         return true;
0219     }
0220
0221     static handle cast(Type &&src, return_value_policy policy, handle parent) {
0222         if (policy == return_value_policy::reference
0223             || policy == return_value_policy::reference_internal) {
0224             pybind11_fail("Cannot use a reference return value policy for an rvalue");
0225         }
0226         return cast_impl(&src, return_value_policy::move, parent);
0227     }
0228
0229     static handle cast(const Type &&src, return_value_policy policy, handle parent) {
0230         if (policy == return_value_policy::reference
0231             || policy == return_value_policy::reference_internal) {
0232             pybind11_fail("Cannot use a reference return value policy for an rvalue");
0233         }
0234         return cast_impl(&src, return_value_policy::move, parent);
0235     }
0236
0237     static handle cast(Type &src, return_value_policy policy, handle parent) {
0238         if (policy == return_value_policy::automatic
0239             || policy == return_value_policy::automatic_reference) {
0240             policy = return_value_policy::copy;
0241         }
0242         return cast_impl(&src, policy, parent);
0243     }
0244
0245     static handle cast(const Type &src, return_value_policy policy, handle parent) {
0246         if (policy == return_value_policy::automatic
0247             || policy == return_value_policy::automatic_reference) {
0248             policy = return_value_policy::copy;
0249         }
0250         return cast(&src, policy, parent);
0251     }
0252
0253     static handle cast(Type *src, return_value_policy policy, handle parent) {
0254         if (policy == return_value_policy::automatic) {
0255             policy = return_value_policy::take_ownership;
0256         } else if (policy == return_value_policy::automatic_reference) {
0257             policy = return_value_policy::reference;
0258         }
0259         return cast_impl(src, policy, parent);
0260     }
0261
0262     static handle cast(const Type *src, return_value_policy policy, handle parent) {
0263         if (policy == return_value_policy::automatic) {
0264             policy = return_value_policy::take_ownership;
0265         } else if (policy == return_value_policy::automatic_reference) {
0266             policy = return_value_policy::reference;
0267         }
0268         return cast_impl(src, policy, parent);
0269     }
0270
0271     template <typename C>
0272     static handle cast_impl(C *src, return_value_policy policy, handle parent) {
0273         object parent_object;
0274         bool writeable = false;
0275         switch (policy) {
0276             case return_value_policy::move:
0277                 if (std::is_const<C>::value) {
0278                     pybind11_fail("Cannot move from a constant reference");
0279                 }
0280
0281                 src = Helper::alloc(std::move(*src));
0282
0283                 parent_object
0284                     = capsule(src, [](void *ptr) { Helper::free(reinterpret_cast<Type *>(ptr)); });
0285                 writeable = true;
0286                 break;
0287
0288             case return_value_policy::take_ownership:
0289                 if (std::is_const<C>::value) {
0290                     // This cast is ugly, and might be UB in some cases, but we don't have an
0291                     // alternative here as we must free that memory
0292                     Helper::free(const_cast<Type *>(src));
0293                     pybind11_fail("Cannot take ownership of a const reference");
0294                 }
0295
0296                 parent_object
0297                     = capsule(src, [](void *ptr) { Helper::free(reinterpret_cast<Type *>(ptr)); });
0298                 writeable = true;
0299                 break;
0300
0301             case return_value_policy::copy:
0302                 writeable = true;
0303                 break;
0304
0305             case return_value_policy::reference:
0306                 parent_object = none();
0307                 writeable = !std::is_const<C>::value;
0308                 break;
0309
0310             case return_value_policy::reference_internal:
0311                 // Default should do the right thing
0312                 if (!parent) {
0313                     pybind11_fail("Cannot use reference internal when there is no parent");
0314                 }
0315                 parent_object = reinterpret_borrow<object>(parent);
0316                 writeable = !std::is_const<C>::value;
0317                 break;
0318
0319             default:
0320                 pybind11_fail("pybind11 bug in eigen.h, please file a bug report");
0321         }
0322
0323         auto result = array_t<typename Type::Scalar, compute_array_flag_from_tensor<Type>()>(
0324             convert_dsizes_to_vector(Helper::get_shape(*src)), src->data(), parent_object);
0325
0326         if (!writeable) {
0327             array_proxy(result.ptr())->flags &= ~detail::npy_api::NPY_ARRAY_WRITEABLE_;
0328         }
0329
0330         return result.release();
0331     }
0332 };
0333
0334 template <typename StoragePointerType,
0335           bool needs_writeable,
0336           enable_if_t<!needs_writeable, bool> = true>
0337 StoragePointerType get_array_data_for_type(array &arr) {
0338 #if EIGEN_VERSION_AT_LEAST(3, 4, 0)
0339     return reinterpret_cast<StoragePointerType>(arr.data());
0340 #else
0341     // Handle Eigen bug
0342     return reinterpret_cast<StoragePointerType>(const_cast<void *>(arr.data()));
0343 #endif
0344 }
0345
0346 template <typename StoragePointerType,
0347           bool needs_writeable,
0348           enable_if_t<needs_writeable, bool> = true>
0349 StoragePointerType get_array_data_for_type(array &arr) {
0350     return reinterpret_cast<StoragePointerType>(arr.mutable_data());
0351 }
0352
0353 template <typename T, typename = void>
0354 struct get_storage_pointer_type;
0355
0356 template <typename MapType>
0357 struct get_storage_pointer_type<MapType, void_t<typename MapType::StoragePointerType>> {
0358     using SPT = typename MapType::StoragePointerType;
0359 };
0360
0361 template <typename MapType>
0362 struct get_storage_pointer_type<MapType, void_t<typename MapType::PointerArgType>> {
0363     using SPT = typename MapType::PointerArgType;
0364 };
0365
0366 template <typename Type, int Options>
0367 struct type_caster<Eigen::TensorMap<Type, Options>,
0368                    typename eigen_tensor_helper<remove_cv_t<Type>>::ValidType> {
0369     static_assert(!std::is_pointer<typename Type::Scalar>::value,
0370                   PYBIND11_EIGEN_MESSAGE_POINTER_TYPES_ARE_NOT_SUPPORTED);
0371     using MapType = Eigen::TensorMap<Type, Options>;
0372     using Helper = eigen_tensor_helper<remove_cv_t<Type>>;
0373
0374     bool load(handle src, bool /*convert*/) {
0375         // Note that we have a lot more checks here as we want to make sure to avoid copies
0376         if (!isinstance<array>(src)) {
0377             return false;
0378         }
0379         auto arr = reinterpret_borrow<array>(src);
0380         if ((arr.flags() & compute_array_flag_from_tensor<Type>()) == 0) {
0381             return false;
0382         }
0383
0384         if (!arr.dtype().is(dtype::of<typename Type::Scalar>())) {
0385             return false;
0386         }
0387
0388         if (arr.ndim() != Type::NumIndices) {
0389             return false;
0390         }
0391
0392         constexpr bool is_aligned = (Options & Eigen::Aligned) != 0;
0393
0394         if (is_aligned && !is_tensor_aligned(arr.data())) {
0395             return false;
0396         }
0397
0398         auto shape = get_shape_for_array<typename Type::Index, Type::NumIndices>(arr);
0399
0400         if (!Helper::is_correct_shape(shape)) {
0401             return false;
0402         }
0403
0404         if (needs_writeable && !arr.writeable()) {
0405             return false;
0406         }
0407
0408         auto result = get_array_data_for_type<typename get_storage_pointer_type<MapType>::SPT,
0409                                               needs_writeable>(arr);
0410
0411         value.reset(new MapType(std::move(result), std::move(shape)));
0412
0413         return true;
0414     }
0415
0416     static handle cast(MapType &&src, return_value_policy policy, handle parent) {
0417         return cast_impl(&src, policy, parent);
0418     }
0419
0420     static handle cast(const MapType &&src, return_value_policy policy, handle parent) {
0421         return cast_impl(&src, policy, parent);
0422     }
0423
0424     static handle cast(MapType &src, return_value_policy policy, handle parent) {
0425         if (policy == return_value_policy::automatic
0426             || policy == return_value_policy::automatic_reference) {
0427             policy = return_value_policy::copy;
0428         }
0429         return cast_impl(&src, policy, parent);
0430     }
0431
0432     static handle cast(const MapType &src, return_value_policy policy, handle parent) {
0433         if (policy == return_value_policy::automatic
0434             || policy == return_value_policy::automatic_reference) {
0435             policy = return_value_policy::copy;
0436         }
0437         return cast(&src, policy, parent);
0438     }
0439
0440     static handle cast(MapType *src, return_value_policy policy, handle parent) {
0441         if (policy == return_value_policy::automatic) {
0442             policy = return_value_policy::take_ownership;
0443         } else if (policy == return_value_policy::automatic_reference) {
0444             policy = return_value_policy::reference;
0445         }
0446         return cast_impl(src, policy, parent);
0447     }
0448
0449     static handle cast(const MapType *src, return_value_policy policy, handle parent) {
0450         if (policy == return_value_policy::automatic) {
0451             policy = return_value_policy::take_ownership;
0452         } else if (policy == return_value_policy::automatic_reference) {
0453             policy = return_value_policy::reference;
0454         }
0455         return cast_impl(src, policy, parent);
0456     }
0457
0458     template <typename C>
0459     static handle cast_impl(C *src, return_value_policy policy, handle parent) {
0460         object parent_object;
0461         constexpr bool writeable = !std::is_const<C>::value;
0462         switch (policy) {
0463             case return_value_policy::reference:
0464                 parent_object = none();
0465                 break;
0466
0467             case return_value_policy::reference_internal:
0468                 // Default should do the right thing
0469                 if (!parent) {
0470                     pybind11_fail("Cannot use reference internal when there is no parent");
0471                 }
0472                 parent_object = reinterpret_borrow<object>(parent);
0473                 break;
0474
0475             default:
0476                 // move, take_ownership don't make any sense for a ref/map:
0477                 pybind11_fail("Invalid return_value_policy for Eigen Map type, must be either "
0478                               "reference or reference_internal");
0479         }
0480
0481         auto result = array_t<typename Type::Scalar, compute_array_flag_from_tensor<Type>()>(
0482             convert_dsizes_to_vector(Helper::get_shape(*src)),
0483             src->data(),
0484             std::move(parent_object));
0485
0486         if (!writeable) {
0487             array_proxy(result.ptr())->flags &= ~detail::npy_api::NPY_ARRAY_WRITEABLE_;
0488         }
0489
0490         return result.release();
0491     }
0492
0493 #if EIGEN_VERSION_AT_LEAST(3, 4, 0)
0494
0495     static constexpr bool needs_writeable = !std::is_const<typename std::remove_pointer<
0496         typename get_storage_pointer_type<MapType>::SPT>::type>::value;
0497 #else
0498     // Handle Eigen bug
0499     static constexpr bool needs_writeable = !std::is_const<Type>::value;
0500 #endif
0501
0502 protected:
0503     // TODO: Move to std::optional once std::optional has more support
0504     std::unique_ptr<MapType> value;
0505
0506 public:
0507     // return_descr forces the use of NDArray instead of ArrayLike since refs can only reference
0508     // arrays
0509     static constexpr auto name
0510         = return_descr(get_tensor_descriptor<Type, true, needs_writeable>::value);
0511     explicit operator MapType *() { return value.get(); }
0512     explicit operator MapType &() { return *value; }
0513     explicit operator MapType &&() && { return std::move(*value); }
0514
0515     template <typename T_>
0516     using cast_op_type = ::pybind11::detail::movable_cast_op_type<T_>;
0517 };
0518
0519 PYBIND11_NAMESPACE_END(detail)
0520 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)