API: Simplify basic iterator constructors, add 'itershape' to advanced iterator constructor

11 files changed, 269 insertions, 116 deletions

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index cd897e3ce..1bb1ce666 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -259,61 +259,62 @@ multiarray_funcs_api = {
     # New Iterator API
     'NpyIter_New':                          224,
     'NpyIter_MultiNew':                     225,
-    'NpyIter_Copy':                         226,
-    'NpyIter_Deallocate':                   227,
-    'NpyIter_HasDelayedBufAlloc':           228,
-    'NpyIter_HasInnerLoop':                 229,
-    'NpyIter_RemoveInnerLoop':              230,
-    'NpyIter_GetInnerStrideArray':          231,
-    'NpyIter_GetInnerLoopSizePtr':          232,
-    'NpyIter_Reset':                        233,
-    'NpyIter_ResetBasePointers':            234,
-    'NpyIter_ResetToIterIndexRange':        235,
-    'NpyIter_GetNDim':                      236,
-    'NpyIter_GetNIter':                     237,
-    'NpyIter_GetIterNext':                  238,
-    'NpyIter_GetIterSize':                  239,
-    'NpyIter_GetIterIndexRange':            240,
-    'NpyIter_GetIterIndex':                 241,
-    'NpyIter_GotoIterIndex':                242,
-    'NpyIter_HasCoords':                    243,
-    'NpyIter_GetShape':                     244,
-    'NpyIter_GetGetCoords':                 245,
-    'NpyIter_GotoCoords':                   246,
-    'NpyIter_RemoveCoords':                 247,
-    'NpyIter_HasIndex':                     248,
-    'NpyIter_IsBuffered':                   249,
-    'NpyIter_IsGrowInner':                  250,
-    'NpyIter_GetBufferSize':                251,
-    'NpyIter_GetIndexPtr':                  252,
-    'NpyIter_GotoIndex':                    253,
-    'NpyIter_GetDataPtrArray':              254,
-    'NpyIter_GetDescrArray':                255,
-    'NpyIter_GetOperandArray':              256,
-    'NpyIter_GetIterView':                  257,
-    'NpyIter_GetReadFlags':                 258,
-    'NpyIter_GetWriteFlags':                259,
-    'NpyIter_DebugPrint':                   260,
-    'NpyIter_IterationNeedsAPI':            261,
-    'NpyIter_GetInnerFixedStrideArray':     262,
-    'NpyIter_RemoveAxis':                   263,
-    'NpyIter_GetAxisStrideArray':           264,
-    'NpyIter_RequiresBuffering':            265,
-    'NpyIter_GetInitialDataPtrArray':       266,
-    'NpyIter_CreateCompatibleStrides':      267,
+    'NpyIter_AdvancedNew':                  226,
+    'NpyIter_Copy':                         227,
+    'NpyIter_Deallocate':                   228,
+    'NpyIter_HasDelayedBufAlloc':           229,
+    'NpyIter_HasInnerLoop':                 230,
+    'NpyIter_RemoveInnerLoop':              231,
+    'NpyIter_GetInnerStrideArray':          232,
+    'NpyIter_GetInnerLoopSizePtr':          233,
+    'NpyIter_Reset':                        234,
+    'NpyIter_ResetBasePointers':            235,
+    'NpyIter_ResetToIterIndexRange':        236,
+    'NpyIter_GetNDim':                      237,
+    'NpyIter_GetNIter':                     238,
+    'NpyIter_GetIterNext':                  239,
+    'NpyIter_GetIterSize':                  240,
+    'NpyIter_GetIterIndexRange':            241,
+    'NpyIter_GetIterIndex':                 242,
+    'NpyIter_GotoIterIndex':                243,
+    'NpyIter_HasCoords':                    244,
+    'NpyIter_GetShape':                     245,
+    'NpyIter_GetGetCoords':                 246,
+    'NpyIter_GotoCoords':                   247,
+    'NpyIter_RemoveCoords':                 248,
+    'NpyIter_HasIndex':                     249,
+    'NpyIter_IsBuffered':                   250,
+    'NpyIter_IsGrowInner':                  251,
+    'NpyIter_GetBufferSize':                252,
+    'NpyIter_GetIndexPtr':                  253,
+    'NpyIter_GotoIndex':                    254,
+    'NpyIter_GetDataPtrArray':              255,
+    'NpyIter_GetDescrArray':                256,
+    'NpyIter_GetOperandArray':              257,
+    'NpyIter_GetIterView':                  258,
+    'NpyIter_GetReadFlags':                 259,
+    'NpyIter_GetWriteFlags':                260,
+    'NpyIter_DebugPrint':                   261,
+    'NpyIter_IterationNeedsAPI':            262,
+    'NpyIter_GetInnerFixedStrideArray':     263,
+    'NpyIter_RemoveAxis':                   264,
+    'NpyIter_GetAxisStrideArray':           265,
+    'NpyIter_RequiresBuffering':            266,
+    'NpyIter_GetInitialDataPtrArray':       267,
+    'NpyIter_CreateCompatibleStrides':      268,
     #
-    'PyArray_CastingConverter':             268,
-    'PyArray_CountNonzero':                 269,
-    'PyArray_PromoteTypes':                 270,
-    'PyArray_MinScalarType':                271,
-    'PyArray_ResultType':                   272,
-    'PyArray_CanCastArrayTo':               273,
-    'PyArray_CanCastTypeTo':                274,
-    'PyArray_EinsteinSum':                  275,
-    'PyArray_NewLikeArray':                 276,
-    'PyArray_GetArrayParamsFromObject':     277,
-    'PyArray_ConvertClipmodeSequence':      278,
-    'PyArray_MatrixProduct2':               279,
+    'PyArray_CastingConverter':             269,
+    'PyArray_CountNonzero':                 270,
+    'PyArray_PromoteTypes':                 271,
+    'PyArray_MinScalarType':                272,
+    'PyArray_ResultType':                   273,
+    'PyArray_CanCastArrayTo':               274,
+    'PyArray_CanCastTypeTo':                275,
+    'PyArray_EinsteinSum':                  276,
+    'PyArray_NewLikeArray':                 277,
+    'PyArray_GetArrayParamsFromObject':     278,
+    'PyArray_ConvertClipmodeSequence':      279,
+    'PyArray_MatrixProduct2':               280,
 }
 
 ufunc_types_api = {
diff --git a/numpy/core/src/multiarray/convert.c b/numpy/core/src/multiarray/convert.c
index b40519a4e..44aff3a8d 100644
--- a/numpy/core/src/multiarray/convert.c
+++ b/numpy/core/src/multiarray/convert.c
@@ -392,8 +392,7 @@ PyArray_FillWithZero(PyArrayObject *a)
 
     /* Use an iterator to go through all the data */
     iter = NpyIter_New(a, NPY_ITER_WRITEONLY|NPY_ITER_NO_INNER_ITERATION,
-                    NPY_KEEPORDER, NPY_NO_CASTING,
-                    NULL, 0, NULL, 0);
+                    NPY_KEEPORDER, NPY_NO_CASTING, NULL);
 
     if (iter == NULL) {
         return -1;
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index 08909819d..369cc411e 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -2486,7 +2486,7 @@ PyArray_CopyAnyIntoOrdered(PyArrayObject *dst, PyArrayObject *src,
                                 NPY_ITER_REFS_OK,
                                 order,
                                 NPY_NO_CASTING,
-                                NULL, 0, NULL, 0);
+                                NULL);
     if (dst_iter == NULL) {
         return -1;
     }
@@ -2496,7 +2496,7 @@ PyArray_CopyAnyIntoOrdered(PyArrayObject *dst, PyArrayObject *src,
                                 NPY_ITER_REFS_OK,
                                 order,
                                 NPY_NO_CASTING,
-                                NULL, 0, NULL, 0);
+                                NULL);
     if (src_iter == NULL) {
         NpyIter_Deallocate(dst_iter);
         return -1;
@@ -2718,7 +2718,7 @@ PyArray_CopyInto(PyArrayObject *dst, PyArrayObject *src)
                             NPY_KEEPORDER,
                             NPY_NO_CASTING,
                             op_flags,
-                            NULL, 0, NULL, 0);
+                            NULL);
         if (iter == NULL) {
             return -1;
         }
diff --git a/numpy/core/src/multiarray/einsum.c.src b/numpy/core/src/multiarray/einsum.c.src
index ab939f734..10c6bdca5 100644
--- a/numpy/core/src/multiarray/einsum.c.src
+++ b/numpy/core/src/multiarray/einsum.c.src
@@ -2965,7 +2965,7 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
                     NPY_ITER_NO_BROADCAST;
 
     /* Allocate the iterator */
-    iter = NpyIter_MultiNew(nop+1, op, NPY_ITER_NO_INNER_ITERATION|
+    iter = NpyIter_AdvancedNew(nop+1, op, NPY_ITER_NO_INNER_ITERATION|
                 ((dtype != NULL) ? 0 : NPY_ITER_COMMON_DTYPE)|
                                        NPY_ITER_BUFFERED|
                                        NPY_ITER_DELAY_BUFALLOC|
@@ -2975,7 +2975,7 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
                                        NPY_ITER_ZEROSIZE_OK,
                                        order, casting,
                                        op_flags, op_dtypes,
-                                       ndim_iter, op_axes, 0);
+                                       ndim_iter, op_axes, NULL, 0);
 
     if (iter == NULL) {
         goto fail;
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
index 4d1621534..c497845a1 100644
--- a/numpy/core/src/multiarray/item_selection.c
+++ b/numpy/core/src/multiarray/item_selection.c
@@ -1729,7 +1729,7 @@ PyArray_CountNonzero(PyArrayObject *self)
                              NPY_ITER_NO_INNER_ITERATION|
                              NPY_ITER_REFS_OK,
                         NPY_KEEPORDER, NPY_NO_CASTING,
-                        NULL, 0, NULL, 0);
+                        NULL);
     if (iter == NULL) {
         return -1;
     }
@@ -1823,7 +1823,7 @@ PyArray_Nonzero(PyArrayObject *self)
                              NPY_ITER_ZEROSIZE_OK|
                              NPY_ITER_REFS_OK,
                         NPY_CORDER, NPY_NO_CASTING,
-                        NULL, 0, NULL, 0);
+                        NULL);
 
     if (iter == NULL) {
         Py_DECREF(ret);
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 833c969b0..d39fdb0a8 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -745,7 +745,7 @@ PyArray_InnerProduct(PyObject *op1, PyObject *op2)
  * just like inner product but does the swapaxes stuff on the fly
  */
 NPY_NO_EXPORT PyObject *
-PyArray_MatrixProduct2(PyObject *op1, PyObject *op2, PyObject* out)
+PyArray_MatrixProduct2(PyObject *op1, PyObject *op2, PyArrayObject* out)
 {
     PyArrayObject *ap1, *ap2, *ret = NULL;
     PyArrayIterObject *it1, *it2;
@@ -2032,7 +2032,12 @@ array_matrixproduct(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject* kwds)
     if (o == Py_None) {
         o = NULL;
     }
-    return _ARET(PyArray_MatrixProduct2(a, v, o));
+    if (o != NULL && !PyArray_Check(o)) {
+        PyErr_SetString(PyExc_TypeError,
+                        "'out' must be an array");
+        return NULL;
+    }
+    return _ARET(PyArray_MatrixProduct2(a, v, (PyArrayObject *)o));
 }
 
 static int
diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index d80fc6edc..559a6f2e5 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -285,7 +285,8 @@ struct NpyIter_AD {
 static int
 npyiter_check_global_flags(npy_uint32 flags, npy_uint32* itflags);
 static int
-npyiter_check_op_axes(int niter, int oa_ndim, int **op_axes);
+npyiter_check_op_axes(int niter, int oa_ndim, int **op_axes,
+                        npy_intp *itershape);
 static int
 npyiter_calculate_ndim(int niter, PyArrayObject **op_in,
                        int oa_ndim);
@@ -315,6 +316,7 @@ static int
 npyiter_fill_axisdata(NpyIter *iter, npy_uint32 flags, char *op_itflags,
                     char **op_dataptr,
                     npy_uint32 *op_flags, int **op_axes,
+                    npy_intp *itershape,
                     int output_scalars);
 static void
 npyiter_replace_axisdata(NpyIter *iter, int iiter,
@@ -372,14 +374,16 @@ npyiter_checkreducesize(NpyIter *iter, npy_intp count,
                                 npy_intp *reduce_outerdim);
 
 /*NUMPY_API
- * Allocate a new iterator for multiple array objects
+ * Allocate a new iterator for multiple array objects, and advanced
+ * options for controlling the broadcasting, shape, and buffer size.
  */
 NPY_NO_EXPORT NpyIter *
-NpyIter_MultiNew(int niter, PyArrayObject **op_in, npy_uint32 flags,
+NpyIter_AdvancedNew(int niter, PyArrayObject **op_in, npy_uint32 flags,
                  NPY_ORDER order, NPY_CASTING casting,
                  npy_uint32 *op_flags,
                  PyArray_Descr **op_request_dtypes,
-                 int oa_ndim, int **op_axes, npy_intp buffersize)
+                 int oa_ndim, int **op_axes, npy_intp *itershape,
+                 npy_intp buffersize)
 {
     npy_uint32 itflags = NPY_ITFLAG_IDENTPERM;
     int idim, ndim;
@@ -433,7 +437,7 @@ NpyIter_MultiNew(int niter, PyArrayObject **op_in, npy_uint32 flags,
     }
 
     /* Error check 'oa_ndim' and 'op_axes', which must be used together */
-    if (!npyiter_check_op_axes(niter, oa_ndim, op_axes)) {
+    if (!npyiter_check_op_axes(niter, oa_ndim, op_axes, itershape)) {
         return NULL;
     }
 
@@ -502,7 +506,8 @@ NpyIter_MultiNew(int niter, PyArrayObject **op_in, npy_uint32 flags,
 
     /* Fill in the AXISDATA arrays and set the ITERSIZE field */
     if (!npyiter_fill_axisdata(iter, flags, op_itflags, op_dataptr,
-                                        op_flags, op_axes, output_scalars)) {
+                                        op_flags, op_axes, itershape,
+                                        output_scalars)) {
         NpyIter_Deallocate(iter);
         return NULL;
     }
@@ -516,7 +521,7 @@ NpyIter_MultiNew(int niter, PyArrayObject **op_in, npy_uint32 flags,
          * small enough to be cache-friendly.
          */
         if (buffersize <= 0) {
-            buffersize = 1 << 12;
+            buffersize = NPY_BUFSIZE;
         }
         /* No point in a buffer bigger than the iteration size */
         if (buffersize > NIT_ITERSIZE(iter)) {
@@ -768,28 +773,35 @@ NpyIter_MultiNew(int niter, PyArrayObject **op_in, npy_uint32 flags,
 }
 
 /*NUMPY_API
- * Allocate a new iterator for one array object
+ * Allocate a new iterator for more than one array object, using
+ * standard NumPy broadcasting rules and the default buffer size.
+ */
+NPY_NO_EXPORT NpyIter *
+NpyIter_MultiNew(int niter, PyArrayObject **op_in, npy_uint32 flags,
+                 NPY_ORDER order, NPY_CASTING casting,
+                 npy_uint32 *op_flags,
+                 PyArray_Descr **op_request_dtypes)
+{
+    return NpyIter_AdvancedNew(niter, op_in, flags, order, casting,
+                            op_flags, op_request_dtypes,
+                            0, NULL, NULL, 0);
+}
+
+/*NUMPY_API
+ * Allocate a new iterator for one array object.
  */
 NPY_NO_EXPORT NpyIter *
 NpyIter_New(PyArrayObject *op, npy_uint32 flags,
                   NPY_ORDER order, NPY_CASTING casting,
-                  PyArray_Descr* dtype,
-                  int a_ndim, int *axes, npy_intp buffersize)
+                  PyArray_Descr* dtype)
 {
     /* Split the flags into separate global and op flags */
     npy_uint32 op_flags = flags&NPY_ITER_PER_OP_FLAGS;
     flags &= NPY_ITER_GLOBAL_FLAGS;
 
-    if (a_ndim > 0) {
-        return NpyIter_MultiNew(1, &op, flags, order, casting,
-                                &op_flags, &dtype,
-                                a_ndim, &axes, buffersize);
-    }
-    else {
-        return NpyIter_MultiNew(1, &op, flags, order, casting,
-                                &op_flags, &dtype,
-                                0, NULL, buffersize);
-    }
+    return NpyIter_AdvancedNew(1, &op, flags, order, casting,
+                            &op_flags, &dtype,
+                            0, NULL, NULL, 0);
 }
 
 /*NUMPY_API
@@ -2903,15 +2915,16 @@ npyiter_calculate_ndim(int niter, PyArrayObject **op_in,
 }
 
 static int
-npyiter_check_op_axes(int niter, int oa_ndim, int **op_axes)
+npyiter_check_op_axes(int niter, int oa_ndim, int **op_axes,
+                        npy_intp *itershape)
 {
     char axes_dupcheck[NPY_MAXDIMS];
     int iiter, idim;
 
-    if (oa_ndim == 0 && op_axes != NULL) {
+    if (oa_ndim == 0 && (op_axes != NULL || itershape != NULL)) {
         PyErr_Format(PyExc_ValueError,
-                "If 'op_axes' is not NULL in the iterator constructor, "
-                "'oa_ndim' must be greater than zero");
+                "If 'op_axes' or 'itershape' is not NULL in the"
+                "iterator constructor, 'oa_ndim' must be greater than zero");
         return 0;
     }
     else if (oa_ndim > 0) {
@@ -3420,6 +3433,7 @@ static int
 npyiter_fill_axisdata(NpyIter *iter, npy_uint32 flags, char *op_itflags,
                     char **op_dataptr,
                     npy_uint32 *op_flags, int **op_axes,
+                    npy_intp *itershape,
                     int output_scalars)
 {
     npy_uint32 itflags = NIT_ITFLAGS(iter);
@@ -3433,8 +3447,19 @@ npyiter_fill_axisdata(NpyIter *iter, npy_uint32 flags, char *op_itflags,
     npy_intp broadcast_shape[NPY_MAXDIMS];
 
     /* First broadcast the shapes together */
-    for (idim = 0; idim < ndim; ++idim) {
-        broadcast_shape[idim] = 1;
+    if (itershape == NULL) {
+        for (idim = 0; idim < ndim; ++idim) {
+            broadcast_shape[idim] = 1;
+        }
+    }
+    else {
+        for (idim = 0; idim < ndim; ++idim) {
+            broadcast_shape[idim] = itershape[idim];
+            /* Negative shape entries are deduced from the operands */
+            if (broadcast_shape[idim] < 0) {
+                broadcast_shape[idim] = 1;
+            }
+        }
     }
     for (iiter = 0; iiter < niter; ++iiter) {
         op_cur = op[iiter];
@@ -3493,6 +3518,17 @@ npyiter_fill_axisdata(NpyIter *iter, npy_uint32 flags, char *op_itflags,
             }
         }
     }
+    /*
+     * If a shape was provided with a 1 entry, make sure that entry didn't
+     * get expanded by broadcasting.
+     */
+    if (itershape != NULL) {
+        for (idim = 0; idim < ndim; ++idim) {
+            if (itershape[idim] == 1 && broadcast_shape[idim] != 1) {
+                goto broadcast_error;
+            }
+        }
+    }
 
     axisdata = NIT_AXISDATA(iter);
     sizeof_axisdata = NIT_AXISDATA_SIZEOF(itflags, ndim, niter);
@@ -3636,12 +3672,16 @@ broadcast_error: {
         if (op_axes == NULL) {
             errmsg = PyUString_FromString("operands could not be broadcast "
                                           "together with shapes ");
+            if (errmsg == NULL) {
+                return 0;
+            }
             for (iiter = 0; iiter < niter; ++iiter) {
                 if (op[iiter] != NULL) {
                     tmp = npyiter_shape_string(PyArray_NDIM(op[iiter]),
                                                     PyArray_DIMS(op[iiter]),
                                                     " ");
                     if (tmp == NULL) {
+                        Py_DECREF(errmsg);
                         return 0;
                     }
                     PyUString_ConcatAndDel(&errmsg, tmp);
@@ -3650,6 +3690,28 @@ broadcast_error: {
                     }
                 }
             }
+            if (itershape != NULL) {
+                tmp = PyUString_FromString("and requested shape ");
+                if (tmp == NULL) {
+                    Py_DECREF(errmsg);
+                    return 0;
+                }
+                PyUString_ConcatAndDel(&errmsg, tmp);
+                if (errmsg == NULL) {
+                    return 0;
+                }
+
+                tmp = npyiter_shape_string(ndim, itershape, "");
+                if (tmp == NULL) {
+                    Py_DECREF(errmsg);
+                    return 0;
+                }
+                PyUString_ConcatAndDel(&errmsg, tmp);
+                if (errmsg == NULL) {
+                    return 0;
+                }
+
+            }
             PyErr_SetObject(PyExc_ValueError, errmsg);
         }
         else {
@@ -3694,6 +3756,28 @@ broadcast_error: {
                     }
                 }
             }
+            if (itershape != NULL) {
+                tmp = PyUString_FromString("and requested shape ");
+                if (tmp == NULL) {
+                    Py_DECREF(errmsg);
+                    return 0;
+                }
+                PyUString_ConcatAndDel(&errmsg, tmp);
+                if (errmsg == NULL) {
+                    return 0;
+                }
+
+                tmp = npyiter_shape_string(ndim, itershape, "");
+                if (tmp == NULL) {
+                    Py_DECREF(errmsg);
+                    return 0;
+                }
+                PyUString_ConcatAndDel(&errmsg, tmp);
+                if (errmsg == NULL) {
+                    return 0;
+                }
+
+            }
             PyErr_SetObject(PyExc_ValueError, errmsg);
         }
 
diff --git a/numpy/core/src/multiarray/new_iterator_pywrap.c b/numpy/core/src/multiarray/new_iterator_pywrap.c
index 8e3231d40..a2af3030b 100644
--- a/numpy/core/src/multiarray/new_iterator_pywrap.c
+++ b/numpy/core/src/multiarray/new_iterator_pywrap.c
@@ -633,7 +633,7 @@ npyiter_convert_op_axes(PyObject *op_axes_in, npy_intp niter,
 }
 
 /*
- * Converts the operand array and op_flags array into the form NpyIter_MultiNew
+ * Converts the operand array and op_flags array into the form NpyIter_AdvancedNew
  * needs.  Sets niter, and on success, each op[i] owns a reference
  * to an array object.
  */
@@ -704,6 +704,7 @@ npyiter_convert_ops(PyObject *op_in, PyObject *op_flags_in,
         for (iiter = 0; iiter < niter; ++iiter) {
             Py_XDECREF(op[iiter]);
         }
+        *niter_out = 0;
         return 0;
     }
 
@@ -729,6 +730,7 @@ npyiter_convert_ops(PyObject *op_in, PyObject *op_flags_in,
                 for (iiter = 0; iiter < niter; ++iiter) {
                     Py_DECREF(op[iiter]);
                 }
+                *niter_out = 0;
                 return 0;
             }
             Py_DECREF(op[iiter]);
@@ -743,7 +745,8 @@ static int
 npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
 {
     static char *kwlist[] = {"op", "flags", "op_flags", "op_dtypes",
-                             "order", "casting", "op_axes", "buffersize",
+                             "order", "casting", "op_axes", "itershape",
+                             "buffersize",
                              NULL};
 
     PyObject *op_in = NULL, *op_flags_in = NULL,
@@ -759,6 +762,7 @@ npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
     int oa_ndim = 0;
     int op_axes_arrays[NPY_MAXARGS][NPY_MAXDIMS];
     int *op_axes[NPY_MAXARGS];
+    PyArray_Dims itershape = {NULL, 0};
     int buffersize = 0;
 
     if (self->iter != NULL) {
@@ -767,7 +771,7 @@ npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
         return -1;
     }
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O&OOO&O&Oi", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O&OOO&O&OO&i", kwlist,
                     &op_in,
                     NpyIter_GlobalFlagsConverter, &flags,
                     &op_flags_in,
@@ -775,19 +779,23 @@ npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
                     npyiter_order_converter, &order,
                     PyArray_CastingConverter, &casting,
                     &op_axes_in,
+                    PyArray_IntpConverter, &itershape,
                     &buffersize)) {
+        if (itershape.ptr != NULL) {
+            PyDimMem_FREE(itershape.ptr);
+        }
         return -1;
     }
 
+    /* Set the dtypes and ops to all NULL to start */
+    memset(op_request_dtypes, 0, sizeof(op_request_dtypes));
+
     /* op and op_flags */
     if (npyiter_convert_ops(op_in, op_flags_in, op, op_flags, &niter)
                                                         != 1) {
-        return -1;
+        goto fail;
     }
 
-    /* Set the dtypes to all NULL to start as well */
-    memset(op_request_dtypes, 0, sizeof(op_request_dtypes[0])*niter);
-
     /* op_request_dtypes */
     if (op_dtypes_in != NULL && op_dtypes_in != Py_None &&
             npyiter_convert_dtypes(op_dtypes_in,
@@ -808,9 +816,27 @@ npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
         }
     }
 
-    self->iter = NpyIter_MultiNew(niter, op, flags, order, casting, op_flags,
+    if (itershape.len > 0) {
+        if (oa_ndim == 0) {
+            oa_ndim = itershape.len;
+            memset(op_axes, 0, sizeof(op_axes[0])*oa_ndim);
+        }
+        else if (oa_ndim != itershape.len) {
+            PyErr_SetString(PyExc_ValueError,
+                        "'op_axes' and 'itershape' must have the same number "
+                        "of entries equal to the iterator ndim");
+            goto fail;
+        }
+    }
+    else if (itershape.ptr != NULL) {
+        PyDimMem_FREE(itershape.ptr);
+        itershape.ptr = NULL;
+    }
+
+    self->iter = NpyIter_AdvancedNew(niter, op, flags, order, casting, op_flags,
                                   op_request_dtypes,
                                   oa_ndim, oa_ndim > 0 ? op_axes : NULL,
+                                  itershape.ptr,
                                   buffersize);
 
     if (self->iter == NULL) {
@@ -829,6 +855,10 @@ npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
         self->finished = 0;
     }
 
+    if (itershape.ptr != NULL) {
+        PyDimMem_FREE(itershape.ptr);
+    }
+
     /* Release the references we got to the ops and dtypes */
     for (iiter = 0; iiter < niter; ++iiter) {
         Py_XDECREF(op[iiter]);
@@ -838,6 +868,9 @@ npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
     return 0;
 
 fail:
+    if (itershape.ptr != NULL) {
+        PyDimMem_FREE(itershape.ptr);
+    }
     for (iiter = 0; iiter < niter; ++iiter) {
         Py_XDECREF(op[iiter]);
         Py_XDECREF(op_request_dtypes[iiter]);
@@ -1059,16 +1092,18 @@ NpyIter_NestedIters(PyObject *NPY_UNUSED(self),
         }
 
         if (inest < nnest-1) {
-            iter->iter = NpyIter_MultiNew(niter, op, flags, order,
+            iter->iter = NpyIter_AdvancedNew(niter, op, flags, order,
                                 casting, op_flags, op_request_dtypes,
                                 nested_naxes[inest], op_axes_niter,
+                                NULL,
                                 0);
         }
         else {
-            iter->iter = NpyIter_MultiNew(niter, op, flags_inner, order,
+            iter->iter = NpyIter_AdvancedNew(niter, op, flags_inner, order,
                                 casting, op_flags_inner,
                                 op_request_dtypes_inner,
                                 nested_naxes[inest], op_axes_niter,
+                                NULL,
                                 buffersize);
         }
 
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index b35b57d86..638032bab 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -1796,7 +1796,7 @@ iterator_loop(PyUFuncObject *self,
      * were already checked, we use the casting rule 'unsafe' which
      * is faster to calculate.
      */
-    iter = NpyIter_MultiNew(niter, op,
+    iter = NpyIter_AdvancedNew(niter, op,
                         NPY_ITER_NO_INNER_ITERATION|
                         NPY_ITER_REFS_OK|
                         NPY_ITER_ZEROSIZE_OK|
@@ -1805,7 +1805,7 @@ iterator_loop(PyUFuncObject *self,
                         NPY_ITER_DELAY_BUFALLOC,
                         order, NPY_UNSAFE_CASTING,
                         op_flags, dtype,
-                        0, NULL, buffersize);
+                        0, NULL, NULL, buffersize);
     if (iter == NULL) {
         return -1;
     }
@@ -2305,11 +2305,11 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *self,
     }
 
     /* Create the iterator */
-    iter = NpyIter_MultiNew(niter, op, NPY_ITER_COORDS|
+    iter = NpyIter_AdvancedNew(niter, op, NPY_ITER_COORDS|
                                       NPY_ITER_REFS_OK|
                                       NPY_ITER_REDUCE_OK,
                            order, NPY_UNSAFE_CASTING, op_flags,
-                           dtype, op_ndim, op_axes, 0);
+                           dtype, op_ndim, op_axes, NULL, 0);
     if (iter == NULL) {
         retval = -1;
         goto fail;
@@ -2924,11 +2924,11 @@ PyUFunc_ReductionOp(PyUFuncObject *self, PyArrayObject *arr,
             op_dtypes[1] = op_dtypes[0];
         }
         NPY_UF_DBG_PRINT("Allocating outer iterator\n");
-        iter = NpyIter_MultiNew(2, op, flags,
+        iter = NpyIter_AdvancedNew(2, op, flags,
                                    NPY_KEEPORDER, NPY_UNSAFE_CASTING,
                                    op_flags,
                                    op_dtypes_param,
-                                   ndim_iter, op_axes, 0);
+                                   ndim_iter, op_axes, NULL, 0);
         if (iter == NULL) {
             goto fail;
         }
@@ -3036,7 +3036,7 @@ PyUFunc_ReductionOp(PyUFuncObject *self, PyArrayObject *arr,
             op_axes[0][0] = -1;
             op_axes[1][0] = axis;
 
-            iter_inner = NpyIter_MultiNew(2, op, NPY_ITER_NO_INNER_ITERATION|
+            iter_inner = NpyIter_AdvancedNew(2, op, NPY_ITER_NO_INNER_ITERATION|
                                        NPY_ITER_BUFFERED|
                                        NPY_ITER_DELAY_BUFALLOC|
                                        NPY_ITER_GROWINNER|
@@ -3044,7 +3044,7 @@ PyUFunc_ReductionOp(PyUFuncObject *self, PyArrayObject *arr,
                                        NPY_ITER_REFS_OK,
                                        NPY_CORDER, NPY_UNSAFE_CASTING,
                                        op_flags, op_dtypes,
-                                       1, op_axes, buffersize);
+                                       1, op_axes, NULL, buffersize);
         }
         /* Should never get an inner iterator for ACCUMULATE */
         else {
@@ -3557,11 +3557,11 @@ PyUFunc_Reduceat(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *ind,
         op_dtypes[1] = op_dtypes[0];
 
         NPY_UF_DBG_PRINT("Allocating outer iterator\n");
-        iter = NpyIter_MultiNew(3, op, flags,
+        iter = NpyIter_AdvancedNew(3, op, flags,
                                    NPY_KEEPORDER, NPY_UNSAFE_CASTING,
                                    op_flags,
                                    op_dtypes,
-                                   ndim, op_axes, 0);
+                                   ndim, op_axes, NULL, 0);
         if (iter == NULL) {
             goto fail;
         }
diff --git a/numpy/core/tests/test_new_iterator.py b/numpy/core/tests/test_new_iterator.py
index 404735223..28cac8563 100644
--- a/numpy/core/tests/test_new_iterator.py
+++ b/numpy/core/tests/test_new_iterator.py
@@ -570,6 +570,35 @@ def test_iter_broadcasting():
     assert_equal(i.itersize, 24)
     assert_equal(i.shape, (4,2,3))
 
+def test_iter_itershape():
+    # Check that allocated outputs work with a specified shape
+    a = np.arange(6, dtype='i2').reshape(2,3)
+    i = newiter([a, None], [], [['readonly'], ['writeonly','allocate']],
+                            op_axes=[[0,1,None], None],
+                            itershape=(-1,-1,4))
+    assert_equal(i.operands[1].shape, (2,3,4))
+    assert_equal(i.operands[1].strides, (24,8,2))
+
+    i = newiter([a.T, None], [], [['readonly'], ['writeonly','allocate']],
+                            op_axes=[[0,1,None], None],
+                            itershape=(-1,-1,4))
+    assert_equal(i.operands[1].shape, (3,2,4))
+    assert_equal(i.operands[1].strides, (8,24,2))
+
+    i = newiter([a.T, None], [], [['readonly'], ['writeonly','allocate']],
+                            order='F',
+                            op_axes=[[0,1,None], None],
+                            itershape=(-1,-1,4))
+    assert_equal(i.operands[1].shape, (3,2,4))
+    assert_equal(i.operands[1].strides, (2,6,12))
+
+    # If we specify 1 in the itershape, it shouldn't allow broadcasting
+    # of that dimension to a bigger value
+    assert_raises(ValueError, newiter, [a, None], [],
+                            [['readonly'], ['writeonly','allocate']],
+                            op_axes=[[0,1,None], None],
+                            itershape=(-1,1,4))
+
 def test_iter_broadcasting_errors():
     # Check that errors are thrown for bad broadcasting shapes
 
diff --git a/numpy/lib/src/_compiled_base.c b/numpy/lib/src/_compiled_base.c
index b44cfb471..d1b2b2c6a 100644
--- a/numpy/lib/src/_compiled_base.c
+++ b/numpy/lib/src/_compiled_base.c
@@ -679,7 +679,8 @@ static PyObject *
 arr_ravel_coords(PyObject *self, PyObject *args, PyObject *kwds)
 {
     int i, s;
-    PyObject *mode0=NULL, *coords0=NULL, *ret;
+    PyObject *mode0=NULL, *coords0=NULL;
+    PyArrayObject *ret = NULL;
     PyArray_Dims dimensions={0,0};
     npy_intp ravel_strides[NPY_MAXDIMS];
     PyArray_ORDER order = NPY_CORDER;
@@ -758,8 +759,7 @@ arr_ravel_coords(PyObject *self, PyObject *args, PyObject *kwds)
                                                   NPY_ITER_ZEROSIZE_OK,
                                                   NPY_KEEPORDER,
                                                   NPY_SAME_KIND_CASTING,
-                                                  op_flags, dtype,
-                                                  0, NULL, 0);
+                                                  op_flags, dtype);
     if (iter == NULL) {
         goto fail;
     }
@@ -787,7 +787,7 @@ arr_ravel_coords(PyObject *self, PyObject *args, PyObject *kwds)
         } while(iternext(iter));
     }
 
-    ret = (PyObject *)NpyIter_GetOperandArray(iter)[dimensions.len];
+    ret = NpyIter_GetOperandArray(iter)[dimensions.len];
     Py_INCREF(ret);
 
     Py_DECREF(dtype[0]);
@@ -925,7 +925,7 @@ arr_unravel_index(PyObject *self, PyObject *args, PyObject *kwds)
                                 NPY_ITER_DONT_NEGATE_STRIDES|
                                 NPY_ITER_COORDS,
                                 NPY_KEEPORDER, NPY_SAME_KIND_CASTING,
-                                dtype, 0, NULL, 0);
+                                dtype);
     if (iter == NULL) {
         goto fail;
     }