From 163d3c533d37fa76727ecdee429446ad41aac7f2 Mon Sep 17 00:00:00 2001 From: mattip Date: Tue, 7 Feb 2023 10:51:58 +0200 Subject: BUG: improve check for 1d operands (from review) --- numpy/core/src/umath/ufunc_object.c | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) (limited to 'numpy/core/src') diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index 63994c258..60efdce96 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -6436,12 +6436,15 @@ ufunc_at(PyUFuncObject *ufunc, PyObject *args) /* * Try to use trivial loop (1d, no casting, aligned) if * - the matching info has a indexed loop + * - idx must be exactly one integer index array * - all operands are 1d + * A future enhancement could loosen the restriction on 1d operands + * by adding an iteration loop inside trivial_at_loop */ if ((ufuncimpl->contiguous_indexed_loop != NULL) && (PyArray_NDIM(op1_array) == 1) && (op2_array != NULL && PyArray_NDIM(op2_array) == 1) && - (iter->nd == 1)) { + (iter->subspace_iter == NULL) && (iter->numiter == 1)) { res = trivial_at_loop(ufuncimpl, flags, iter, op1_array, op2_array, &context); -- cgit v1.2.1