Repository: numpy/numpy
File: numpy/_core/src/multiarray/mapping.c (C)

Code before:
        }

        /* shape is set, and strides is set up to mit->nd, set rest */
        PyArray_CreateSortedStridePerm(PyArray_NDIM(subspace),
                                PyArray_STRIDES(subspace), strideperm);
        stride = extra_op_dtype->elsize;
        for (i=PyArray_NDIM(subspace) - 1; i >= 0; i--) {
            strides[mit->nd_fancy + strideperm[i].perm] = stride;
            stride *= PyArray_DIM(subspace, (int)strideperm[i].perm);
        }

        /*
         * Allocate new array. Note: Always base class, because
         * subclasses might mess with the shape.
         */
        Py_INCREF(extra_op_dtype);
        extra_op = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type,
                                           extra_op_dtype,
                                           mit->nd_fancy + PyArray_NDIM(subspace),
                                           mit->dimensions, strides,
                                           NULL, 0, NULL);
        if (extra_op == NULL) {
            goto fail;
        }
    }

    /*
     * The extra op is now either allocated, can be allocated by
     * NpyIter (no subspace) or is not used at all.
     *
     * Need to set the axis remapping for the extra_op. This needs
     * to cause ignoring of subspace dimensions and prepending -1
     * for broadcasting.
     */
    if (extra_op) {
        for (j=0; j < mit->nd - PyArray_NDIM(extra_op); j++) {
            single_op_axis[j] = -1;
        }
        for (i=0; i < PyArray_NDIM(extra_op); i++) {
            /* (fills subspace dimensions too, but they are not unused) */
            single_op_axis[j++] = i;
        }
    }

    /*
     * NOTE: If for some reason someone wishes to use REDUCE_OK, be
     *       careful and fix the error message replacement at the end.
     */
    outer_flags = NPY_ITER_ZEROSIZE_OK |
                  NPY_ITER_REFS_OK |
                  NPY_ITER_BUFFERED |

Review comment (suggestion):
LGTM, but the ref-counting looks off.  We own a reference to `extra_op_dtype` but incref it above for cleanup.

I would be OK with just using `Py_SETREF(extra_op_dtype, Py_NewRef(...))` though.  That way we just always own one reference.

Code after:
        }

        /* shape is set, and strides is set up to mit->nd, set rest */
        PyArray_CreateSortedStridePerm(PyArray_NDIM(subspace),
                                PyArray_STRIDES(subspace), strideperm);
        stride = extra_op_dtype->elsize;
        for (i=PyArray_NDIM(subspace) - 1; i >= 0; i--) {
            strides[mit->nd_fancy + strideperm[i].perm] = stride;
            stride *= PyArray_DIM(subspace, (int)strideperm[i].perm);
        }

        /*
         * Allocate new array. Note: Always base class, because
         * subclasses might mess with the shape.
         */
        Py_INCREF(extra_op_dtype);
        extra_op = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type,
                                           extra_op_dtype,
                                           mit->nd_fancy + PyArray_NDIM(subspace),
                                           mit->dimensions, strides,
                                           NULL, 0, NULL);
        if (extra_op == NULL) {
            goto fail;
        }
        // extra_op_dtype might have been replaced, so get a new reference
        extra_op_dtype = PyArray_DESCR(extra_op);
    }

    /*
     * The extra op is now either allocated, can be allocated by
     * NpyIter (no subspace) or is not used at all.
     *
     * Need to set the axis remapping for the extra_op. This needs
     * to cause ignoring of subspace dimensions and prepending -1
     * for broadcasting.
     */
    if (extra_op) {
        for (j=0; j < mit->nd - PyArray_NDIM(extra_op); j++) {
            single_op_axis[j] = -1;
        }
        for (i=0; i < PyArray_NDIM(extra_op); i++) {
            /* (fills subspace dimensions too, but they are not unused) */
            single_op_axis[j++] = i;
        }
    }

    /*
     * NOTE: If for some reason someone wishes to use REDUCE_OK, be
     *       careful and fix the error message replacement at the end.
     */
    outer_flags = NPY_ITER_ZEROSIZE_OK |