ENH: add np.resize and np.ndarray.resize

The former fills with repeated copies of input, while the latter
fills with zeros.

Author: ev-br

torch_np/_funcs.py

@@ -18,6 +18,7 @@
 # these implement ndarray methods but need not be public functions
 semi_private = [
     "_flatten",
+    "_ndarray_resize",
 ]
 
 

torch_np/_funcs_impl.py

@@ -6,6 +6,8 @@
 # Contents of this module ends up in the main namespace via _funcs.py
 # where type annotations are used in conjunction with the @normalizer decorator.
 
+import builtins
+import math
 import operator
 from typing import Optional, Sequence
 
@@ -43,7 +45,6 @@ def copyto(dst: NDArray, src: ArrayLike, casting="same_kind", where=NoValue):
     dst.tensor.copy_(src)
 
 
-
 def atleast_1d(*arys: ArrayLike):
     res = torch.atleast_1d(*arys)
     if isinstance(res, tuple):
@@ -995,6 +996,65 @@ def tile(A: ArrayLike, reps):
     return torch.tile(A, reps)
 
 
+def resize(a: ArrayLike, new_shape=None):
+    # implementation vendored from
+    # https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/fromnumeric.py#L1420-L1497
+    if new_shape is None:
+        return a
+
+    if isinstance(new_shape, int):
+        new_shape = (new_shape,)
+
+    a = ravel(a)
+
+    new_size = 1
+    for dim_length in new_shape:
+        new_size *= dim_length
+        if dim_length < 0:
+            raise ValueError("all elements of `new_shape` must be non-negative")
+
+    if a.numel() == 0 or new_size == 0:
+        # First case must zero fill. The second would have repeats == 0.
+        return torch.zeros(new_shape, dtype=a.dtype)
+
+    repeats = -(-new_size // a.numel())  # ceil division
+    a = concatenate((a,) * repeats)[:new_size]
+
+    return reshape(a, new_shape)
+
+
+def _ndarray_resize(a: ArrayLike, new_shape, refcheck=False):
+    # implementation of ndarray.resize.
+    # NB: differs from np.resize: fills with zeros instead of making repeated copies of input.
+    if refcheck:
+        raise NotImplementedError(
+            f"resize(..., refcheck={refcheck} is not implemented."
+        )
+
+    if new_shape in [(), (None,)]:
+        return a
+
+    # support both x.resize((2, 2)) and x.resize(2, 2)
+    if len(new_shape) == 1:
+        new_shape = new_shape[0]
+    if isinstance(new_shape, int):
+        new_shape = (new_shape,)
+
+    a = ravel(a)
+
+    if builtins.any(x < 0 for x in new_shape):
+        raise ValueError("all elements of `new_shape` must be non-negative")
+
+    new_numel = math.prod(new_shape)
+    if new_numel < a.numel():
+        # shrink
+        return a[:new_numel].reshape(new_shape)
+    else:
+        b = torch.zeros(new_numel)
+        b[: a.numel()] = a
+        return b.reshape(new_shape)
+
+
 # ### diag et al ###
 
 
@@ -1821,7 +1881,6 @@ def common_type(*tensors: ArrayLike):
         return array_type[0][precision]
 
 
-
 # ### histograms ###
 
 
@@ -1864,4 +1923,3 @@ def histogram(
         b = b.to(int)
 
     return h, b
-

torch_np/_ndarray.py

@@ -358,6 +358,12 @@ def reshape(self, *shape, order="C"):
     ravel = _funcs.ravel
     flatten = _funcs._flatten
 
+    def resize(self, *new_shape, refcheck=False):
+        # ndarray.resize works in-place (may cause a reallocation though)
+        self.tensor = _funcs_impl._ndarray_resize(
+            self.tensor, new_shape, refcheck=refcheck
+        )
+
     nonzero = _funcs.nonzero
     clip = _funcs.clip
     repeat = _funcs.repeat

torch_np/tests/numpy_tests/core/test_multiarray.py

@@ -4513,7 +4513,6 @@ def test_index_getset(self):
         assert it.index == it.base.size
 
 
-@pytest.mark.xfail(reason='TODO')
 class TestResize:
 
     @_no_tracing
@@ -4523,10 +4522,11 @@ def test_basic(self):
             x.resize((5, 5), refcheck=False)
         else:
             x.resize((5, 5))
-        assert_array_equal(x.flat[:9],
-                np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).flat)
-        assert_array_equal(x[9:].flat, 0)
+        assert_array_equal(x.ravel()[:9],
+                np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).ravel())
+        assert_array_equal(x[9:].ravel(), 0)
 
+    @pytest.mark.skip(reason="how to find if someone is refencing an array")
     def test_check_reference(self):
         x = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
         y = x
@@ -4565,7 +4565,7 @@ def test_invalid_arguments(self):
         assert_raises(TypeError, np.eye(3).resize, 'hi')
         assert_raises(ValueError, np.eye(3).resize, -1)
         assert_raises(TypeError, np.eye(3).resize, order=1)
-        assert_raises(TypeError, np.eye(3).resize, refcheck='hi')
+        assert_raises((NotImplementedError, TypeError), np.eye(3).resize, refcheck='hi')
 
     @_no_tracing
     def test_freeform_shape(self):
@@ -4586,18 +4586,6 @@ def test_zeros_appended(self):
         assert_array_equal(x[0], np.eye(3))
         assert_array_equal(x[1], np.zeros((3, 3)))
 
-    @_no_tracing
-    def test_obj_obj(self):
-        # check memory is initialized on resize, gh-4857
-        a = np.ones(10, dtype=[('k', object, 2)])
-        if IS_PYPY:
-            a.resize(15, refcheck=False)
-        else:
-            a.resize(15,)
-        assert_equal(a.shape, (15,))
-        assert_array_equal(a['k'][-5:], 0)
-        assert_array_equal(a['k'][:-5], 1)
-
     def test_empty_view(self):
         # check that sizes containing a zero don't trigger a reallocate for
         # already empty arrays
@@ -4606,6 +4594,7 @@ def test_empty_view(self):
         x_view.resize((0, 10))
         x_view.resize((0, 100))
 
+    @pytest.mark.skip(reason="ignore weakrefs for ndarray.resize")
     def test_check_weakref(self):
         x = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
         xref = weakref.ref(x)

torch_np/tests/numpy_tests/core/test_numeric.py

@@ -26,7 +26,6 @@
 from hypothesis.extra import numpy as hynp
 
 
-@pytest.mark.xfail(reason="TODO")
 class TestResize:
     def test_copies(self):
         A = np.array([[1, 2], [3, 4]])
@@ -64,28 +63,17 @@ def test_zeroresize(self):
 
     def test_reshape_from_zero(self):
         # See also gh-6740
-        A = np.zeros(0, dtype=[('a', np.float32)])
+        A = np.zeros(0, dtype=np.float32)
         Ar = np.resize(A, (2, 1))
         assert_array_equal(Ar, np.zeros((2, 1), Ar.dtype))
         assert_equal(A.dtype, Ar.dtype)
 
     def test_negative_resize(self):
         A = np.arange(0, 10, dtype=np.float32)
         new_shape = (-10, -1)
-        with pytest.raises(ValueError, match=r"negative"):
+        with pytest.raises((RuntimeError, ValueError)):
             np.resize(A, new_shape=new_shape)
 
-    def test_subclass(self):
-        class MyArray(np.ndarray):
-            __array_priority__ = 1.
-
-        my_arr = np.array([1]).view(MyArray)
-        assert type(np.resize(my_arr, 5)) is MyArray
-        assert type(np.resize(my_arr, 0)) is MyArray
-
-        my_arr = np.array([]).view(MyArray)
-        assert type(np.resize(my_arr, 5)) is MyArray
-
 
 class TestNonarrayArgs:
     # check that non-array arguments to functions wrap them in arrays