8000 Fix index_tricks issue by 87 · Pull Request #445 · numpy/numpy · GitHub
[go: up one dir, main page]

Skip to content

Fix index_tricks issue #445

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
BUG: Fix problems with ndindex and nditer
This fixes an issue with ndindex shape tuple recognition, and an issue
in the nditer where scalar input did not produce an empty index tuple.
To be able to fix nditer, an extra flag has been added: NPY_ITFLAG_SCALAR
and a new function NpyIter_IsScalar has been added to the nditer API.
Also a few tests have been added to make sure the ndindex behaves as
intended.
  • Loading branch information
87 committed Sep 19, 2012
commit 3043864aac4942957a3e44a8c915220709da35e4
1 change: 1 addition & 0 deletions numpy/core/code_generators/numpy_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,7 @@
'PyDataMem_FREE': 289,
'PyDataMem_RENEW': 290,
'PyDataMem_SetEventHook': 291,
'NpyIter_IsScalar': 292,
}

ufunc_types_api = {
Expand Down
9 changes: 9 additions & 0 deletions numpy/core/src/multiarray/nditer_api.c
Original file line number Diff line number Diff line change
Expand Up @@ -844,6 +844,15 @@ NpyIter_IsGrowInner(NpyIter *iter)
return (NIT_ITFLAGS(iter)&NPY_ITFLAG_GROWINNER) != 0;
}

/*NUMPY_API
* Whether the iterator output is scalar
*/
NPY_NO_EXPORT npy_bool
NpyIter_IsScalar(NpyIter *iter)
{
return (NIT_ITFLAGS(iter)&NPY_ITFLAG_SCALAR) != 0;
}

/*NUMPY_API
* Gets the size of the buffer, or 0 if buffering is not enabled
*/
Expand Down
32 changes: 13 additions & 19 deletions numpy/core/src/multiarray/nditer_constr.c
Original file line number Diff line number Diff line change
Expand Up @@ -54,8 +54,7 @@ static int
npyiter_fill_axisdata(NpyIter *iter, npy_uint32 flags, npyiter_opitflags *op_itflags,
char **op_dataptr,
npy_uint32 *op_flags, int **op_axes,
npy_intp *itershape,
int output_scalars);
npy_intp *itershape);
static void
npyiter_replace_axisdata(NpyIter *iter, int iop,
PyArrayObject *op,
Expand All @@ -74,8 +73,7 @@ npyiter_find_best_axis_ordering(NpyIter *iter);
static PyArray_Descr *
npyiter_get_common_dtype(int nop, PyArrayObject **op,
npyiter_opitflags *op_itflags, PyArray_Descr **op_dtype,
PyArray_Descr **op_request_dtypes,
int only_inputs, int output_scalars);
PyArray_Descr **op_request_dtypes, int only_inputs);
static PyArrayObject *
npyiter_new_temp_array(NpyIter *iter, PyTypeObject *subtype,
npy_uint32 flags, npyiter_opitflags *op_itflags,
Expand All @@ -86,7 +84,7 @@ npyiter_allocate_arrays(NpyIter *iter,
npy_uint32 flags,
PyArray_Descr **op_dtype, PyTypeObject *subtype,
npy_uint32 *op_flags, npyiter_opitflags *op_itflags,
int **op_axes, int output_scalars);
int **op_axes);
static void
npyiter_get_priority_subtype(int nop, PyArrayObject **op,
npyiter_opitflags *op_itflags,
Expand Down Expand Up @@ -123,7 +121,7 @@ NpyIter_AdvancedNew(int nop, PyArrayObject **op_in, npy_uint32 flags,
npy_int8 *perm;
NpyIter_BufferData *bufferdata = NULL;
int any_allocate = 0, any_missing_dtypes = 0,
output_scalars = 0, need_subtype = 0;
need_subtype = 0;

/* The subtype for automatically allocated outputs */
double subtype_priority = NPY_PRIORITY;
Expand Down Expand Up @@ -177,7 +175,7 @@ NpyIter_AdvancedNew(int nop, PyArrayObject **op_in, npy_uint32 flags,

/* If 'ndim' is zero, any outputs should be scalars */
if (ndim == 0) {
output_scalars = 1;
itflags |= NPY_ITFLAG_SCALAR;
ndim = 1;
}

Expand Down Expand Up @@ -231,8 +229,7 @@ NpyIter_AdvancedNew(int nop, 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, itershape,
output_scalars)) {
op_flags, op_axes, itershape)) {
NpyIter_Deallocate(iter);
return NULL;
}
Expand Down Expand Up @@ -338,8 +335,7 @@ NpyIter_AdvancedNew(int nop, PyArrayObject **op_in, npy_uint32 flags,
dtype = npyiter_get_common_dtype(nop, op,
op_itflags, op_dtype,
op_request_dtypes,
only_inputs,
output_scalars);
only_inputs);
if (dtype == NULL) {
NpyIter_Deallocate(iter);
return NULL;
Expand Down Expand Up @@ -389,7 +385,7 @@ NpyIter_AdvancedNew(int nop, PyArrayObject **op_in, npy_uint32 flags,
* done now using a memory layout matching the iterator.
*/
if (!npyiter_allocate_arrays(iter, flags, op_dtype, subtype, op_flags,
op_itflags, op_axes, output_scalars)) {
op_itflags, op_axes)) {
NpyIter_Deallocate(iter);
return NULL;
}
Expand Down Expand Up @@ -1437,8 +1433,7 @@ static int
npyiter_fill_axisdata(NpyIter *iter, npy_uint32 flags, npyiter_opitflags *op_itflags,
char **op_dataptr,
npy_uint32 *op_flags, int **op_axes,
npy_intp *itershape,
int output_scalars)
npy_intp *itershape)
{
npy_uint32 itflags = NIT_ITFLAGS(iter);
int idim, ndim = NIT_NDIM(iter);
Expand Down Expand Up @@ -1558,7 +1553,7 @@ npyiter_fill_axisdata(NpyIter *iter, npy_uint32 flags, npyiter_opitflags *op_itf
ondim = PyArray_NDIM(op_cur);
if (bshape == 1) {
strides[iop] = 0;
if (idim >= ondim && !output_scalars &&
if (idim >= ondim && !(itflags & NPY_ITFLAG_SCALAR) &&
(op_flags[iop] & NPY_ITER_NO_BROADCAST)) {
goto operand_different_than_broadcast;
}
Expand Down Expand Up @@ -2393,8 +2388,7 @@ npyiter_find_best_axis_ordering(NpyIter *iter)
static PyArray_Descr *
npyiter_get_common_dtype(int nop, PyArrayObject **op,
npyiter_opitflags *op_itflags, PyArray_Descr **op_dtype,
PyArray_Descr **op_request_dtypes,
int only_inputs, int output_scalars)
PyArray_Descr **op_request_dtypes, int only_inputs)
{
int iop;
npy_intp narrs = 0, ndtypes = 0;
Expand Down Expand Up @@ -2693,7 +2687,7 @@ npyiter_allocate_arrays(NpyIter *iter,
npy_uint32 flags,
PyArray_Descr **op_dtype, PyTypeObject *subtype,
npy_uint32 *op_flags, npyiter_opitflags *op_itflags,
int **op_axes, int output_scalars)
int **op_axes)
{
npy_uint32 itflags = NIT_ITFLAGS(iter);
int idim, ndim = NIT_NDIM(iter);
Expand Down Expand Up @@ -2724,7 +2718,7 @@ npyiter_allocate_arrays(NpyIter *iter,
if (op[iop] == NULL) {
PyArrayObject *out;
PyTypeObject *op_subtype;
int ondim = output_scalars ? 0 : ndim;
int ondim = (itflags & NPY_ITFLAG_SCALAR) ? 0 : ndim;

/* Check whether the subtype was disabled */
op_subtype = (op_flags[iop] & NPY_ITER_NO_SUBTYPE) ?
Expand Down
3 changes: 3 additions & 0 deletions numpy/core/src/multiarray/nditer_impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,8 @@
#define NPY_ITFLAG_REDUCE 0x1000
/* Reduce iteration doesn't need to recalculate reduce loops next time */
#define NPY_ITFLAG_REUSE_REDUCE_LOOPS 0x2000
/* The iterator output is scalar */
#define NPY_ITFLAG_SCALAR 0x4000

/* Internal iterator per-operand iterator flags */

Expand Down Expand Up @@ -215,6 +217,7 @@ typedef npy_int16 npyiter_opitflags;
&(iter)->iter_flexdata + NIT_RESETDATAPTR_OFFSET(itflags, ndim, nop)))
#define NIT_BASEOFFSETS(iter) ((npy_intp *)( \
&(iter)->iter_flexdata + NIT_BASEOFFSETS_OFFSET(itflags, ndim, nop)))

#define NIT_OPERANDS(iter) ((PyArrayObject **)( \
&(iter)->iter_flexdata + NIT_OPERANDS_OFFSET(itflags, ndim, nop)))
#define NIT_OPITFLAGS(iter) ((npyiter_opitflags *)( \
Expand Down
4 changes: 4 additions & 0 deletions numpy/core/src/multiarray/nditer_pywrap.c
Original file line number Diff line number Diff line change
Expand Up @@ -1542,6 +1542,9 @@ static PyObject *npyiter_multi_index_get(NewNpyArrayIterObject *self)
}

if (self->get_multi_index != NULL) {
if (NpyIter_IsScalar(self->iter)) {
return PyTuple_New(0);
}
ndim = NpyIter_GetNDim(self->iter);
self->get_multi_index(self->iter, multi_index);
ret = PyTuple_New(ndim);
Expand Down Expand Up @@ -1968,6 +1971,7 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
return NULL;
}


if (NpyIter_HasDelayedBufAlloc(self->iter)) {
PyErr_SetString(PyExc_ValueError,
"Iterator construction used delayed buffer allocation, "
Expand Down
3 changes: 3 additions & 0 deletions numpy/lib/index_tricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -533,6 +533,9 @@ class ndindex(object):

"""
def __init__(self, *shape):
# Accept shapes in the form f(x, y, ..) as well as f((x ,y, ..))
if len(shape) == 1 and isinstance(shape[0], tuple):
shape = shape[0]
x = as_strided(_nx.zeros(1), shape=shape, strides=_nx.zeros_like(shape))
self._it = _nx.nditer(x, flags=['multi_index'], order='C')

Expand Down
6 changes: 6 additions & 0 deletions numpy/lib/tests/test_index_tricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -241,6 +241,12 @@ def test_ndindex():
x = list(np.ndindex(1, 2, 3))
expected = [ix for ix, e in np.ndenumerate(np.zeros((1, 2, 3)))]
assert_array_equal(x, expected)
# Packed as well as unpacked tuple are acceptable
y = list(np.ndindex((1, 2, 3)))
assert_array_equal(x, y)
# Empty shape gives empty index
z = list(np.ndindex(()))
assert_equal(z, [()])


if __name__ == "__main__":
Expand Down
0