support passing int as shape to broadcast_to

cupy · Dec 30, 2022 · 3efff6b · 3efff6b
1 parent 309d002
commit 3efff6b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/cupy/_core/_routines_manipulation.pyx b/cupy/_core/_routines_manipulation.pyx
@@ -476,6 +476,7 @@ cpdef _ndarray_base broadcast_to(_ndarray_base array, shape):
         :meth:`numpy.broadcast_to`
 
     """
+    shape = tuple(shape) if numpy.iterable(shape) else (shape,)
     cdef int i, j, ndim = array._shape.size(), length = len(shape)
     cdef Py_ssize_t sh, a_sh
     if ndim > length:

diff --git a/tests/cupy_tests/manipulation_tests/test_dims.py b/tests/cupy_tests/manipulation_tests/test_dims.py
@@ -55,6 +55,12 @@ def test_broadcast_to(self, xp, dtype):
         b = xp.broadcast_to(a, (2, 3, 3, 4))
         return b
 
+    @testing.numpy_cupy_array_equal()
+    def test_broadcast_to_int(self, xp):
+        a = testing.shaped_arange((10,), xp, xp.float32)
+        b = xp.broadcast_to(a, 10)
+        return b
+
     @testing.for_all_dtypes()
     def test_broadcast_to_fail(self, dtype):
         for xp in (numpy, cupy):