The source code and dockerfile for the GSW2024 AI Lab.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
This repo is archived. You can view files and clone it, but cannot push or open issues/pull-requests.

300 lines
12 KiB

4 weeks ago
  1. /*
  2. tests/test_numpy_array.cpp -- test core array functionality
  3. Copyright (c) 2016 Ivan Smirnov <i.s.smirnov@gmail.com>
  4. All rights reserved. Use of this source code is governed by a
  5. BSD-style license that can be found in the LICENSE file.
  6. */
  7. #include "pybind11_tests.h"
  8. #include <pybind11/numpy.h>
  9. #include <pybind11/stl.h>
  10. #include <cstdint>
  11. using arr = py::array;
  12. using arr_t = py::array_t<uint16_t, 0>;
  13. static_assert(std::is_same<arr_t::value_type, uint16_t>::value, "");
  14. template<typename... Ix> arr data(const arr& a, Ix... index) {
  15. return arr(a.nbytes() - a.offset_at(index...), (const uint8_t *) a.data(index...));
  16. }
  17. template<typename... Ix> arr data_t(const arr_t& a, Ix... index) {
  18. return arr(a.size() - a.index_at(index...), a.data(index...));
  19. }
  20. arr& mutate_data(arr& a) {
  21. auto ptr = (uint8_t *) a.mutable_data();
  22. for (ssize_t i = 0; i < a.nbytes(); i++)
  23. ptr[i] = (uint8_t) (ptr[i] * 2);
  24. return a;
  25. }
  26. arr_t& mutate_data_t(arr_t& a) {
  27. auto ptr = a.mutable_data();
  28. for (ssize_t i = 0; i < a.size(); i++)
  29. ptr[i]++;
  30. return a;
  31. }
  32. template<typename... Ix> arr& mutate_data(arr& a, Ix... index) {
  33. auto ptr = (uint8_t *) a.mutable_data(index...);
  34. for (ssize_t i = 0; i < a.nbytes() - a.offset_at(index...); i++)
  35. ptr[i] = (uint8_t) (ptr[i] * 2);
  36. return a;
  37. }
  38. template<typename... Ix> arr_t& mutate_data_t(arr_t& a, Ix... index) {
  39. auto ptr = a.mutable_data(index...);
  40. for (ssize_t i = 0; i < a.size() - a.index_at(index...); i++)
  41. ptr[i]++;
  42. return a;
  43. }
  44. template<typename... Ix> ssize_t index_at(const arr& a, Ix... idx) { return a.index_at(idx...); }
  45. template<typename... Ix> ssize_t index_at_t(const arr_t& a, Ix... idx) { return a.index_at(idx...); }
  46. template<typename... Ix> ssize_t offset_at(const arr& a, Ix... idx) { return a.offset_at(idx...); }
  47. template<typename... Ix> ssize_t offset_at_t(const arr_t& a, Ix... idx) { return a.offset_at(idx...); }
  48. template<typename... Ix> ssize_t at_t(const arr_t& a, Ix... idx) { return a.at(idx...); }
  49. template<typename... Ix> arr_t& mutate_at_t(arr_t& a, Ix... idx) { a.mutable_at(idx...)++; return a; }
  50. #define def_index_fn(name, type) \
  51. sm.def(#name, [](type a) { return name(a); }); \
  52. sm.def(#name, [](type a, int i) { return name(a, i); }); \
  53. sm.def(#name, [](type a, int i, int j) { return name(a, i, j); }); \
  54. sm.def(#name, [](type a, int i, int j, int k) { return name(a, i, j, k); });
  55. template <typename T, typename T2> py::handle auxiliaries(T &&r, T2 &&r2) {
  56. if (r.ndim() != 2) throw std::domain_error("error: ndim != 2");
  57. py::list l;
  58. l.append(*r.data(0, 0));
  59. l.append(*r2.mutable_data(0, 0));
  60. l.append(r.data(0, 1) == r2.mutable_data(0, 1));
  61. l.append(r.ndim());
  62. l.append(r.itemsize());
  63. l.append(r.shape(0));
  64. l.append(r.shape(1));
  65. l.append(r.size());
  66. l.append(r.nbytes());
  67. return l.release();
  68. }
  69. test_initializer numpy_array([](py::module &m) {
  70. auto sm = m.def_submodule("array");
  71. sm.def("ndim", [](const arr& a) { return a.ndim(); });
  72. sm.def("shape", [](const arr& a) { return arr(a.ndim(), a.shape()); });
  73. sm.def("shape", [](const arr& a, ssize_t dim) { return a.shape(dim); });
  74. sm.def("strides", [](const arr& a) { return arr(a.ndim(), a.strides()); });
  75. sm.def("strides", [](const arr& a, ssize_t dim) { return a.strides(dim); });
  76. sm.def("writeable", [](const arr& a) { return a.writeable(); });
  77. sm.def("size", [](const arr& a) { return a.size(); });
  78. sm.def("itemsize", [](const arr& a) { return a.itemsize(); });
  79. sm.def("nbytes", [](const arr& a) { return a.nbytes(); });
  80. sm.def("owndata", [](const arr& a) { return a.owndata(); });
  81. def_index_fn(data, const arr&);
  82. def_index_fn(data_t, const arr_t&);
  83. def_index_fn(index_at, const arr&);
  84. def_index_fn(index_at_t, const arr_t&);
  85. def_index_fn(offset_at, const arr&);
  86. def_index_fn(offset_at_t, const arr_t&);
  87. def_index_fn(mutate_data, arr&);
  88. def_index_fn(mutate_data_t, arr_t&);
  89. def_index_fn(at_t, const arr_t&);
  90. def_index_fn(mutate_at_t, arr_t&);
  91. sm.def("make_f_array", [] {
  92. return py::array_t<float>({ 2, 2 }, { 4, 8 });
  93. });
  94. sm.def("make_c_array", [] {
  95. return py::array_t<float>({ 2, 2 }, { 8, 4 });
  96. });
  97. sm.def("wrap", [](py::array a) {
  98. return py::array(
  99. a.dtype(),
  100. {a.shape(), a.shape() + a.ndim()},
  101. {a.strides(), a.strides() + a.ndim()},
  102. a.data(),
  103. a
  104. );
  105. });
  106. struct ArrayClass {
  107. int data[2] = { 1, 2 };
  108. ArrayClass() { py::print("ArrayClass()"); }
  109. ~ArrayClass() { py::print("~ArrayClass()"); }
  110. };
  111. py::class_<ArrayClass>(sm, "ArrayClass")
  112. .def(py::init<>())
  113. .def("numpy_view", [](py::object &obj) {
  114. py::print("ArrayClass::numpy_view()");
  115. ArrayClass &a = obj.cast<ArrayClass&>();
  116. return py::array_t<int>({2}, {4}, a.data, obj);
  117. }
  118. );
  119. sm.def("function_taking_uint64", [](uint64_t) { });
  120. sm.def("isinstance_untyped", [](py::object yes, py::object no) {
  121. return py::isinstance<py::array>(yes) && !py::isinstance<py::array>(no);
  122. });
  123. sm.def("isinstance_typed", [](py::object o) {
  124. return py::isinstance<py::array_t<double>>(o) && !py::isinstance<py::array_t<int>>(o);
  125. });
  126. sm.def("default_constructors", []() {
  127. return py::dict(
  128. "array"_a=py::array(),
  129. "array_t<int32>"_a=py::array_t<std::int32_t>(),
  130. "array_t<double>"_a=py::array_t<double>()
  131. );
  132. });
  133. sm.def("converting_constructors", [](py::object o) {
  134. return py::dict(
  135. "array"_a=py::array(o),
  136. "array_t<int32>"_a=py::array_t<std::int32_t>(o),
  137. "array_t<double>"_a=py::array_t<double>(o)
  138. );
  139. });
  140. // Overload resolution tests:
  141. sm.def("overloaded", [](py::array_t<double>) { return "double"; });
  142. sm.def("overloaded", [](py::array_t<float>) { return "float"; });
  143. sm.def("overloaded", [](py::array_t<int>) { return "int"; });
  144. sm.def("overloaded", [](py::array_t<unsigned short>) { return "unsigned short"; });
  145. sm.def("overloaded", [](py::array_t<long long>) { return "long long"; });
  146. sm.def("overloaded", [](py::array_t<std::complex<double>>) { return "double complex"; });
  147. sm.def("overloaded", [](py::array_t<std::complex<float>>) { return "float complex"; });
  148. sm.def("overloaded2", [](py::array_t<std::complex<double>>) { return "double complex"; });
  149. sm.def("overloaded2", [](py::array_t<double>) { return "double"; });
  150. sm.def("overloaded2", [](py::array_t<std::complex<float>>) { return "float complex"; });
  151. sm.def("overloaded2", [](py::array_t<float>) { return "float"; });
  152. // Only accept the exact types:
  153. sm.def("overloaded3", [](py::array_t<int>) { return "int"; }, py::arg().noconvert());
  154. sm.def("overloaded3", [](py::array_t<double>) { return "double"; }, py::arg().noconvert());
  155. // Make sure we don't do unsafe coercion (e.g. float to int) when not using forcecast, but
  156. // rather that float gets converted via the safe (conversion to double) overload:
  157. sm.def("overloaded4", [](py::array_t<long long, 0>) { return "long long"; });
  158. sm.def("overloaded4", [](py::array_t<double, 0>) { return "double"; });
  159. // But we do allow conversion to int if forcecast is enabled (but only if no overload matches
  160. // without conversion)
  161. sm.def("overloaded5", [](py::array_t<unsigned int>) { return "unsigned int"; });
  162. sm.def("overloaded5", [](py::array_t<double>) { return "double"; });
  163. // Issue 685: ndarray shouldn't go to std::string overload
  164. sm.def("issue685", [](std::string) { return "string"; });
  165. sm.def("issue685", [](py::array) { return "array"; });
  166. sm.def("issue685", [](py::object) { return "other"; });
  167. sm.def("proxy_add2", [](py::array_t<double> a, double v) {
  168. auto r = a.mutable_unchecked<2>();
  169. for (ssize_t i = 0; i < r.shape(0); i++)
  170. for (ssize_t j = 0; j < r.shape(1); j++)
  171. r(i, j) += v;
  172. }, py::arg().noconvert(), py::arg());
  173. sm.def("proxy_init3", [](double start) {
  174. py::array_t<double, py::array::c_style> a({ 3, 3, 3 });
  175. auto r = a.mutable_unchecked<3>();
  176. for (ssize_t i = 0; i < r.shape(0); i++)
  177. for (ssize_t j = 0; j < r.shape(1); j++)
  178. for (ssize_t k = 0; k < r.shape(2); k++)
  179. r(i, j, k) = start++;
  180. return a;
  181. });
  182. sm.def("proxy_init3F", [](double start) {
  183. py::array_t<double, py::array::f_style> a({ 3, 3, 3 });
  184. auto r = a.mutable_unchecked<3>();
  185. for (ssize_t k = 0; k < r.shape(2); k++)
  186. for (ssize_t j = 0; j < r.shape(1); j++)
  187. for (ssize_t i = 0; i < r.shape(0); i++)
  188. r(i, j, k) = start++;
  189. return a;
  190. });
  191. sm.def("proxy_squared_L2_norm", [](py::array_t<double> a) {
  192. auto r = a.unchecked<1>();
  193. double sumsq = 0;
  194. for (ssize_t i = 0; i < r.shape(0); i++)
  195. sumsq += r[i] * r(i); // Either notation works for a 1D array
  196. return sumsq;
  197. });
  198. sm.def("proxy_auxiliaries2", [](py::array_t<double> a) {
  199. auto r = a.unchecked<2>();
  200. auto r2 = a.mutable_unchecked<2>();
  201. return auxiliaries(r, r2);
  202. });
  203. // Same as the above, but without a compile-time dimensions specification:
  204. sm.def("proxy_add2_dyn", [](py::array_t<double> a, double v) {
  205. auto r = a.mutable_unchecked();
  206. if (r.ndim() != 2) throw std::domain_error("error: ndim != 2");
  207. for (ssize_t i = 0; i < r.shape(0); i++)
  208. for (ssize_t j = 0; j < r.shape(1); j++)
  209. r(i, j) += v;
  210. }, py::arg().noconvert(), py::arg());
  211. sm.def("proxy_init3_dyn", [](double start) {
  212. py::array_t<double, py::array::c_style> a({ 3, 3, 3 });
  213. auto r = a.mutable_unchecked();
  214. if (r.ndim() != 3) throw std::domain_error("error: ndim != 3");
  215. for (ssize_t i = 0; i < r.shape(0); i++)
  216. for (ssize_t j = 0; j < r.shape(1); j++)
  217. for (ssize_t k = 0; k < r.shape(2); k++)
  218. r(i, j, k) = start++;
  219. return a;
  220. });
  221. sm.def("proxy_auxiliaries2_dyn", [](py::array_t<double> a) {
  222. return auxiliaries(a.unchecked(), a.mutable_unchecked());
  223. });
  224. sm.def("array_auxiliaries2", [](py::array_t<double> a) {
  225. return auxiliaries(a, a);
  226. });
  227. // Issue #785: Uninformative "Unknown internal error" exception when constructing array from empty object:
  228. sm.def("array_fail_test", []() { return py::array(py::object()); });
  229. sm.def("array_t_fail_test", []() { return py::array_t<double>(py::object()); });
  230. // Make sure the error from numpy is being passed through:
  231. sm.def("array_fail_test_negative_size", []() { int c = 0; return py::array(-1, &c); });
  232. // Issue (unnumbered; reported in #788): regression: initializer lists can be ambiguous
  233. sm.def("array_initializer_list", []() { return py::array_t<float>(1); }); // { 1 } also works, but clang warns about it
  234. sm.def("array_initializer_list", []() { return py::array_t<float>({ 1, 2 }); });
  235. sm.def("array_initializer_list", []() { return py::array_t<float>({ 1, 2, 3 }); });
  236. sm.def("array_initializer_list", []() { return py::array_t<float>({ 1, 2, 3, 4 }); });
  237. // reshape array to 2D without changing size
  238. sm.def("array_reshape2", [](py::array_t<double> a) {
  239. const ssize_t dim_sz = (ssize_t)std::sqrt(a.size());
  240. if (dim_sz * dim_sz != a.size())
  241. throw std::domain_error("array_reshape2: input array total size is not a squared integer");
  242. a.resize({dim_sz, dim_sz});
  243. });
  244. // resize to 3D array with each dimension = N
  245. sm.def("array_resize3", [](py::array_t<double> a, size_t N, bool refcheck) {
  246. a.resize({N, N, N}, refcheck);
  247. });
  248. // return 2D array with Nrows = Ncols = N
  249. sm.def("create_and_resize", [](size_t N) {
  250. py::array_t<double> a;
  251. a.resize({N, N});
  252. std::fill(a.mutable_data(), a.mutable_data() + a.size(), 42.);
  253. return a;
  254. });
  255. });