From fb8fa529a19798941067387d24c3440a1764b125 Mon Sep 17 00:00:00 2001 From: JohnnyOu Date: Fri, 6 May 2022 02:36:26 -0500 Subject: [PATCH] BUG: Fix segmentation fault (#21436) * Fixed #21301 * Fixed #21436, better error message, organized test cases as requested * Update numpy/core/src/umath/ufunc_object.c Co-authored-by: Sebastian Berg --- numpy/core/src/umath/ufunc_object.c | 7 +++++++ numpy/core/tests/test_ufunc.py | 16 +++++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index 1772eb250271..aeca872dd640 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -5884,6 +5884,13 @@ ufunc_at(PyUFuncObject *ufunc, PyObject *args) NPY_BEGIN_THREADS_DEF; + if (ufunc->core_enabled) { + PyErr_Format(PyExc_TypeError, + "%s.at does not support ufunc with non-trivial signature: %s has signature %s.", + ufunc->name, ufunc->name, ufunc->core_signature); + return NULL; + } + if (ufunc->nin > 2) { PyErr_SetString(PyExc_ValueError, "Only unary and binary ufuncs supported at this time"); diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index 6c56e631590e..292797c6d411 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -389,7 +389,7 @@ def test_signature8(self): assert_equal(ixs, (0, 0, 0, 1, 2)) assert_equal(flags, (self.can_ignore, self.size_inferred, 0)) assert_equal(sizes, (3, -1, 9)) - + def test_signature9(self): enabled, num_dims, ixs, flags, sizes = umt.test_signature( 1, 1, "( 3) -> ( )") @@ -2002,6 +2002,20 @@ def test_inplace_fancy_indexing(self): # Test multiple output ufuncs raise error, gh-5665 assert_raises(ValueError, np.modf.at, np.arange(10), [1]) + # Test maximum + a = np.array([1, 2, 3]) + np.maximum.at(a, [0], 0) + assert_equal(np.array([1, 2, 3]), a) + + def test_at_not_none_signature(self): + # Test ufuncs with non-trivial signature raise a TypeError + a = np.ones((2, 2, 2)) + b = np.ones((1, 2, 2)) + assert_raises(TypeError, np.matmul.at, a, [0], b) + + a = np.array([[[1, 2], [3, 4]]]) + assert_raises(TypeError, np.linalg._umath_linalg.det.at, a, [0]) + def test_reduce_arguments(self): f = np.add.reduce d = np.ones((5,2), dtype=int)