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

Code before:
            PyArray_STRIDES(self), dtype->elsize
        );
    }

    nonzero = PyDataType_GetArrFuncs(PyArray_DESCR(self))->nonzero;
    /* If it's a trivial one-dimensional loop, don't use an iterator */
    if (PyArray_TRIVIALLY_ITERABLE(self)) {
        needs_api = PyDataType_FLAGCHK(dtype, NPY_NEEDS_PYAPI);
        PyArray_PREPARE_TRIVIAL_ITERATION(self, count, data, stride);

        if (needs_api){
            while (count--) {
                if (nonzero(data, self)) {
                    ++nonzero_count;
                }
                if (PyErr_Occurred()) {
                    return -1;
                }
                data += stride;
            }
        }
        else {
            /* Special low-overhead version specific to the float types (and some others) */
            npy_intp dispatched_nonzero_count = count_nonzero_trivial_dispatcher(count, data, stride, dtype->type_num);
            if (dispatched_nonzero_count>=0)
                return dispatched_nonzero_count;

            NPY_BEGIN_THREADS_THRESHOLDED(count);
            while (count--) {
                if (nonzero(data, self)) {
                    ++nonzero_count;
                }
                data += stride;
            }
            NPY_END_THREADS;
        }

        return nonzero_count;
    }

    /*
     * If the array has size zero, return zero (the iterator rejects
     * size zero arrays)
     */
    if (PyArray_SIZE(self) == 0) {
        return 0;
    }

    /*
     * Otherwise create and use an iterator to count the nonzeros.
     */

Review comment (suggestion):
```suggestion
            if (dispatched_nonzero_count >= 0) {
                return dispatched_nonzero_count;
            }
```

Code after:
    int needs_api = 0;
    PyArray_Descr *dtype;

    NpyIter *iter;
    NpyIter_IterNextFunc *iternext;
    char **dataptr;
    npy_intp *strideptr, *innersizeptr;
    NPY_BEGIN_THREADS_DEF;

    dtype = PyArray_DESCR(self);
    /* Special low-overhead version specific to the boolean/int types */
    if (PyArray_ISALIGNED(self) && (
            PyDataType_ISBOOL(dtype) || PyDataType_ISINTEGER(dtype))) {
        return count_nonzero_int(
            PyArray_NDIM(self), PyArray_BYTES(self), PyArray_DIMS(self),
            PyArray_STRIDES(self), dtype->elsize
        );
    }

    nonzero = PyDataType_GetArrFuncs(PyArray_DESCR(self))->nonzero;
    /* If it's a trivial one-dimensional loop, don't use an iterator */
    if (PyArray_TRIVIALLY_ITERABLE(self)) {
        needs_api = PyDataType_FLAGCHK(dtype, NPY_NEEDS_PYAPI);
        PyArray_PREPARE_TRIVIAL_ITERATION(self, count, data, stride);

        if (needs_api){
            while (count--) {
                if (nonzero(data, self)) {
                    ++nonzero_count;
                }
                if (PyErr_Occurred()) {
                    return -1;
                }
                data += stride;
            }
        }
        else {
            /* Special low-overhead version specific to the float types (and some others) */
            if (PyArray_ISNOTSWAPPED(self) && PyArray_ISALIGNED(self)) {
                npy_intp dispatched_nonzero_count = count_nonzero_trivial_dispatcher(count,
                                                        data, stride, dtype->type_num);
                if (dispatched_nonzero_count >= 0) {
                    return dispatched_nonzero_count;
                }
            }

            NPY_BEGIN_THREADS_THRESHOLDED(count);
            while (count--) {
                if (nonzero(data, self)) {
                    ++nonzero_count;
                }