334 lines
10 KiB
Plaintext
334 lines
10 KiB
Plaintext
|
/*
|
||
|
|
tests/eigen_tensor.cpp -- automatic conversion of Eigen Tensor
|
||
|
|
|
||
|
|
All rights reserved. Use of this source code is governed by a
|
||
|
|
BSD-style license that can be found in the LICENSE file.
|
||
|
|
*/
|
||
|
|
|
||
|
|
#include <pybind11/eigen/tensor.h>
|
||
|
|
|
||
|
|
PYBIND11_NAMESPACE_BEGIN(eigen_tensor_test)
|
||
|
|
|
||
|
|
namespace py = pybind11;
|
||
|
|
|
||
|
|
PYBIND11_WARNING_DISABLE_MSVC(4127)
|
||
|
|
|
||
|
|
template <typename M>
|
||
|
|
void reset_tensor(M &x) {
|
||
|
|
for (int i = 0; i < x.dimension(0); i++) {
|
||
|
|
for (int j = 0; j < x.dimension(1); j++) {
|
||
|
|
for (int k = 0; k < x.dimension(2); k++) {
|
||
|
|
x(i, j, k) = i * (5 * 2) + j * 2 + k;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
template <typename M>
|
||
|
|
bool check_tensor(M &x) {
|
||
|
|
for (int i = 0; i < x.dimension(0); i++) {
|
||
|
|
for (int j = 0; j < x.dimension(1); j++) {
|
||
|
|
for (int k = 0; k < x.dimension(2); k++) {
|
||
|
|
if (x(i, j, k) != (i * (5 * 2) + j * 2 + k)) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
template <int Options>
|
||
|
|
Eigen::Tensor<double, 3, Options> &get_tensor() {
|
||
|
|
static Eigen::Tensor<double, 3, Options> *x;
|
||
|
|
|
||
|
|
if (!x) {
|
||
|
|
x = new Eigen::Tensor<double, 3, Options>(3, 5, 2);
|
||
|
|
reset_tensor(*x);
|
||
|
|
}
|
||
|
|
|
||
|
|
return *x;
|
||
|
|
}
|
||
|
|
|
||
|
|
template <int Options>
|
||
|
|
Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> &get_tensor_map() {
|
||
|
|
static Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> *x;
|
||
|
|
|
||
|
|
if (!x) {
|
||
|
|
x = new Eigen::TensorMap<Eigen::Tensor<double, 3, Options>>(get_tensor<Options>());
|
||
|
|
}
|
||
|
|
|
||
|
|
return *x;
|
||
|
|
}
|
||
|
|
|
||
|
|
template <int Options>
|
||
|
|
Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options> &get_fixed_tensor() {
|
||
|
|
static Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options> *x;
|
||
|
|
|
||
|
|
if (!x) {
|
||
|
|
Eigen::aligned_allocator<Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>>
|
||
|
|
allocator;
|
||
|
|
x = new (allocator.allocate(1))
|
||
|
|
Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>();
|
||
|
|
reset_tensor(*x);
|
||
|
|
}
|
||
|
|
|
||
|
|
return *x;
|
||
|
|
}
|
||
|
|
|
||
|
|
template <int Options>
|
||
|
|
const Eigen::Tensor<double, 3, Options> &get_const_tensor() {
|
||
|
|
return get_tensor<Options>();
|
||
|
|
}
|
||
|
|
|
||
|
|
template <int Options>
|
||
|
|
struct CustomExample {
|
||
|
|
CustomExample() : member(get_tensor<Options>()), view_member(member) {}
|
||
|
|
|
||
|
|
Eigen::Tensor<double, 3, Options> member;
|
||
|
|
Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> view_member;
|
||
|
|
};
|
||
|
|
|
||
|
|
template <int Options>
|
||
|
|
void init_tensor_module(pybind11::module &m) {
|
||
|
|
const char *needed_options = "";
|
||
|
|
if (Options == Eigen::ColMajor) {
|
||
|
|
needed_options = "F";
|
||
|
|
} else {
|
||
|
|
needed_options = "C";
|
||
|
|
}
|
||
|
|
m.attr("needed_options") = needed_options;
|
||
|
|
|
||
|
|
m.def("setup", []() {
|
||
|
|
reset_tensor(get_tensor<Options>());
|
||
|
|
reset_tensor(get_fixed_tensor<Options>());
|
||
|
|
});
|
||
|
|
|
||
|
|
m.def("is_ok", []() {
|
||
|
|
return check_tensor(get_tensor<Options>()) && check_tensor(get_fixed_tensor<Options>());
|
||
|
|
});
|
||
|
|
|
||
|
|
py::class_<CustomExample<Options>>(m, "CustomExample", py::module_local())
|
||
|
|
.def(py::init<>())
|
||
|
|
.def_readonly(
|
||
|
|
"member", &CustomExample<Options>::member, py::return_value_policy::reference_internal)
|
||
|
|
.def_readonly("member_view",
|
||
|
|
&CustomExample<Options>::view_member,
|
||
|
|
py::return_value_policy::reference_internal);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"copy_fixed_tensor",
|
||
|
|
[]() { return &get_fixed_tensor<Options>(); },
|
||
|
|
py::return_value_policy::copy);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"copy_tensor", []() { return &get_tensor<Options>(); }, py::return_value_policy::copy);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"copy_const_tensor",
|
||
|
|
[]() { return &get_const_tensor<Options>(); },
|
||
|
|
py::return_value_policy::copy);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"move_fixed_tensor_copy",
|
||
|
|
[]() -> Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options> {
|
||
|
|
return get_fixed_tensor<Options>();
|
||
|
|
},
|
||
|
|
py::return_value_policy::move);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"move_tensor_copy",
|
||
|
|
[]() -> Eigen::Tensor<double, 3, Options> { return get_tensor<Options>(); },
|
||
|
|
py::return_value_policy::move);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"move_const_tensor",
|
||
|
|
[]() -> const Eigen::Tensor<double, 3, Options> & { return get_const_tensor<Options>(); },
|
||
|
|
py::return_value_policy::move);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"take_fixed_tensor",
|
||
|
|
|
||
|
|
[]() {
|
||
|
|
Eigen::aligned_allocator<
|
||
|
|
Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>>
|
||
|
|
allocator;
|
||
|
|
return new (allocator.allocate(1))
|
||
|
|
Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>(
|
||
|
|
get_fixed_tensor<Options>());
|
||
|
|
},
|
||
|
|
py::return_value_policy::take_ownership);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"take_tensor",
|
||
|
|
[]() { return new Eigen::Tensor<double, 3, Options>(get_tensor<Options>()); },
|
||
|
|
py::return_value_policy::take_ownership);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"take_const_tensor",
|
||
|
|
[]() -> const Eigen::Tensor<double, 3, Options> * {
|
||
|
|
return new Eigen::Tensor<double, 3, Options>(get_tensor<Options>());
|
||
|
|
},
|
||
|
|
py::return_value_policy::take_ownership);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"take_view_tensor",
|
||
|
|
[]() -> const Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> * {
|
||
|
|
return new Eigen::TensorMap<Eigen::Tensor<double, 3, Options>>(get_tensor<Options>());
|
||
|
|
},
|
||
|
|
py::return_value_policy::take_ownership);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_tensor",
|
||
|
|
[]() { return &get_tensor<Options>(); },
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_tensor_v2",
|
||
|
|
[]() -> Eigen::Tensor<double, 3, Options> & { return get_tensor<Options>(); },
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_tensor_internal",
|
||
|
|
[]() { return &get_tensor<Options>(); },
|
||
|
|
py::return_value_policy::reference_internal);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_fixed_tensor",
|
||
|
|
[]() { return &get_tensor<Options>(); },
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_const_tensor",
|
||
|
|
[]() { return &get_const_tensor<Options>(); },
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_const_tensor_v2",
|
||
|
|
[]() -> const Eigen::Tensor<double, 3, Options> & { return get_const_tensor<Options>(); },
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_view_of_tensor",
|
||
|
|
[]() -> Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> {
|
||
|
|
return get_tensor_map<Options>();
|
||
|
|
},
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_view_of_tensor_v2",
|
||
|
|
// NOLINTNEXTLINE(readability-const-return-type)
|
||
|
|
[]() -> const Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> {
|
||
|
|
return get_tensor_map<Options>(); // NOLINT(readability-const-return-type)
|
||
|
|
}, // NOLINT(readability-const-return-type)
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_view_of_tensor_v3",
|
||
|
|
[]() -> Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> * {
|
||
|
|
return &get_tensor_map<Options>();
|
||
|
|
},
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_view_of_tensor_v4",
|
||
|
|
[]() -> const Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> * {
|
||
|
|
return &get_tensor_map<Options>();
|
||
|
|
},
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_view_of_tensor_v5",
|
||
|
|
[]() -> Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> & {
|
||
|
|
return get_tensor_map<Options>();
|
||
|
|
},
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_view_of_tensor_v6",
|
||
|
|
[]() -> const Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> & {
|
||
|
|
return get_tensor_map<Options>();
|
||
|
|
},
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"reference_view_of_fixed_tensor",
|
||
|
|
[]() {
|
||
|
|
return Eigen::TensorMap<
|
||
|
|
Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>>(
|
||
|
|
get_fixed_tensor<Options>());
|
||
|
|
},
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def("round_trip_tensor",
|
||
|
|
[](const Eigen::Tensor<double, 3, Options> &tensor) { return tensor; });
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"round_trip_tensor_noconvert",
|
||
|
|
[](const Eigen::Tensor<double, 3, Options> &tensor) { return tensor; },
|
||
|
|
py::arg("tensor").noconvert());
|
||
|
|
|
||
|
|
m.def("round_trip_tensor2",
|
||
|
|
[](const Eigen::Tensor<int32_t, 3, Options> &tensor) { return tensor; });
|
||
|
|
|
||
|
|
m.def("round_trip_fixed_tensor",
|
||
|
|
[](const Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options> &tensor) {
|
||
|
|
return tensor;
|
||
|
|
});
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"round_trip_view_tensor",
|
||
|
|
[](Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> view) { return view; },
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"round_trip_view_tensor_ref",
|
||
|
|
[](Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> &view) { return view; },
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"round_trip_view_tensor_ptr",
|
||
|
|
[](Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> *view) { return view; },
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"round_trip_aligned_view_tensor",
|
||
|
|
[](Eigen::TensorMap<Eigen::Tensor<double, 3, Options>, Eigen::Aligned> view) {
|
||
|
|
return view;
|
||
|
|
},
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"round_trip_const_view_tensor",
|
||
|
|
[](Eigen::TensorMap<const Eigen::Tensor<double, 3, Options>> view) {
|
||
|
|
return Eigen::Tensor<double, 3, Options>(view);
|
||
|
|
},
|
||
|
|
py::return_value_policy::move);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"round_trip_rank_0",
|
||
|
|
[](const Eigen::Tensor<double, 0, Options> &tensor) { return tensor; },
|
||
|
|
py::return_value_policy::move);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"round_trip_rank_0_noconvert",
|
||
|
|
[](const Eigen::Tensor<double, 0, Options> &tensor) { return tensor; },
|
||
|
|
py::arg("tensor").noconvert(),
|
||
|
|
py::return_value_policy::move);
|
||
|
|
|
||
|
|
m.def(
|
||
|
|
"round_trip_rank_0_view",
|
||
|
|
[](Eigen::TensorMap<Eigen::Tensor<double, 0, Options>> &tensor) { return tensor; },
|
||
|
|
py::return_value_policy::reference);
|
||
|
|
}
|
||
|
|
|
||
|
|
void test_module(py::module_ &m) {
|
||
|
|
auto f_style = m.def_submodule("f_style");
|
||
|
|
auto c_style = m.def_submodule("c_style");
|
||
|
|
|
||
|
|
init_tensor_module<Eigen::ColMajor>(f_style);
|
||
|
|
init_tensor_module<Eigen::RowMajor>(c_style);
|
||
|
|
}
|
||
|
|
|
||
|
|
PYBIND11_NAMESPACE_END(eigen_tensor_test)
|