reshape for numpy arrays

include/pybind11/numpy.h

@@ -203,6 +203,8 @@ struct npy_api {
     // Unused. Not removed because that affects ABI of the class.
     int (*PyArray_SetBaseObject_)(PyObject *, PyObject *);
     PyObject* (*PyArray_Resize_)(PyObject*, PyArray_Dims*, int, int);
+    PyObject* (*PyArray_Newshape_)(PyObject*, PyArray_Dims*, int);
+
 private:
     enum functions {
         API_PyArray_GetNDArrayCFeatureVersion = 211,
@@ -213,6 +215,7 @@ struct npy_api {
         API_PyArray_DescrFromScalar = 57,
         API_PyArray_FromAny = 69,
         API_PyArray_Resize = 80,
+        API_PyArray_Newshape = 135,
         API_PyArray_CopyInto = 82,
         API_PyArray_NewCopy = 85,
         API_PyArray_NewFromDescr = 94,
@@ -244,6 +247,7 @@ struct npy_api {
         DECL_NPY_API(PyArray_DescrFromScalar);
         DECL_NPY_API(PyArray_FromAny);
         DECL_NPY_API(PyArray_Resize);
+        DECL_NPY_API(PyArray_Newshape);
         DECL_NPY_API(PyArray_CopyInto);
         DECL_NPY_API(PyArray_NewCopy);
         DECL_NPY_API(PyArray_NewFromDescr);
@@ -253,6 +257,7 @@ struct npy_api {
         DECL_NPY_API(PyArray_GetArrayParamsFromObject);
         DECL_NPY_API(PyArray_Squeeze);
         DECL_NPY_API(PyArray_SetBaseObject);
+
 #undef DECL_NPY_API
         return api;
     }
@@ -790,6 +795,15 @@ class array : public buffer {
         if (isinstance<array>(new_array)) { *this = std::move(new_array); }
     }
 
+    array reshape(ShapeContainer new_shape) {
+        detail::npy_api::PyArray_Dims d
+            = {reinterpret_cast<Py_intptr_t *>(new_shape->data()), int(new_shape->size())};
+        // try to reshape, set ordering param to 0 cause it's not used anyway
+        return reinterpret_steal<array>(
+              detail::npy_api::get().PyArray_Newshape_(m_ptr, &d, 0));
+    }
+
+
     /// Ensure that the argument is a NumPy array
     /// In case of an error, nullptr is returned and the Python error is cleared.
     static array ensure(handle h, int ExtraFlags = 0) {

tests/test_numpy_array.cpp

@@ -405,6 +405,19 @@ TEST_SUBMODULE(numpy_array, sm) {
         return a;
     });
 
+    sm.def("array_reshape1", [](py::array_t<double> a, size_t N) {
+        return a.reshape({N, N, N});
+    });
+
+    sm.def("create_and_reshape", [](size_t N, size_t M, size_t O) {
+        py::array_t<double> a;
+        a.resize({N*M*O});
+        std::fill(a.mutable_data(), a.mutable_data() + a.size(), 42.);
+        return a.reshape({N, M, O});
+    });
+    sm.def("reshape_tuple", [](py::array_t<double> a, const std::vector<int> &new_shape) {
+        return a.reshape(new_shape);
+    });
     sm.def("index_using_ellipsis",
            [](const py::array &a) { return a[py::make_tuple(0, py::ellipsis(), 0)]; });
 

tests/test_numpy_array.py

@@ -477,6 +477,24 @@ def test_array_create_and_resize(msg):
     assert np.all(a == 42.0)
 
 
+def test_array_reshape(msg):
+    a = np.random.randn(10 * 10 * 10).astype("float64")
+    x = m.array_reshape1(a, 10)
+    assert x.shape == (10, 10, 10)
+
+
+@pytest.mark.xfail("env.PYPY")
+def test_create_and_reshape(msg):
+    x = m.create_and_reshape(10, 20, 30)
+    assert x.shape == (10, 20, 30)
+
+
+def test_reshape_tuple(msg):
+    a = np.random.randn(10 * 10 * 10).astype("float64")
+    x = m.reshape_tuple(a, (10, 10, 10))
+    assert x.shape == (10, 10, 10)
+
+
 def test_index_using_ellipsis():
     a = m.index_using_ellipsis(np.zeros((5, 6, 7)))
     assert a.shape == (6,)