fix bug with broadcasting to higher dimension

juliantaylor · juliantaylor · commit 94f711de9866 · 2016-12-17T20:31:33.000+01:00
2,1 + 1, -&gt; 2,2 which cannot be done inplace, so check for equal
dimensionality
diff --git a/numpy/core/src/multiarray/number.c b/numpy/core/src/multiarray/number.c
@@ -612,19 +612,14 @@ can_elide_temp(PyArrayObject * alhs, PyObject * orhs, int * cannot)
             return 0;
         }
 
-        /* too large to fit into left hand side */
-        if (PyArray_NDIM(arhs) > PyArray_NDIM(alhs)) {
-            Py_DECREF(arhs);
-            return 0;
-        }
-
         /*
          * if rhs is not a scalar dimensions must match
-         * todo: one could allow full broadcasting on equal types
+         * TODO: one could allow broadcasting on equal types
          */
-        if (PyArray_NDIM(arhs) > 0 &&
-            !PyArray_CompareLists(PyArray_DIMS(alhs), PyArray_DIMS(arhs),
-                                 PyArray_NDIM(arhs))) {
+        if (!(PyArray_NDIM(arhs) == 0 ||
+              (PyArray_NDIM(arhs) == PyArray_NDIM(alhs) &&
+               PyArray_CompareLists(PyArray_DIMS(alhs), PyArray_DIMS(arhs),
+                                    PyArray_NDIM(arhs))))) {
                 Py_DECREF(arhs);
                 return 0;
         }
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
@@ -2741,6 +2741,16 @@ def test_temporary_with_cast(self):
         d = f.astype(np.float64)
         assert_equal(((f + f) + d).dtype, np.dtype('f8'))
 
+    def test_elide_broadcast(self):
+        # test no elision on broadcast to higher dimension
+        # only triggers elision code path in debug mode as triggering it in
+        # normal mode needs 256kb large matching dimension, so a lot of memory
+        d = np.ones((2000, 1), dtype=int)
+        b = np.ones((2000), dtype=np.bool)
+        r = (1 - d) + b
+        assert_equal(r, 1)
+        assert_equal(r.shape, (2000, 2000))
+
     def test_ufunc_override_rop_precedence(self):
         # 2016-01-29: NUMPY_UFUNC_DISABLED
         return
