8000 ENH: add identity kwarg to frompyfunc (#8255) · numpy/numpy@bf6859c · GitHub
[go: up one dir, main page]
Skip to content

Commit bf6859c

Browse files
mattharriganeric-wieser
authored and
seberg
committed
ENH: add identity kwarg to frompyfunc (#8255)
* ENH: add identity kwarg to frompyfunc * Update umathmodule.c * Add test, docs, and release note for identity Co-authored-by: Eric Wieser <wieser.eric@gmail.com>
1 parent 5e57e87 commit bf6859c

File tree

4 files changed

+48
-8
lines changed

4 files changed

+48
-8
lines changed

‎doc/release/upcoming_changes/8255.new_feature.rst

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
`numpy.frompyfunc` now accepts an identity argument
2+
---------------------------------------------------
3+
This allows the :attr:`numpy.ufunc.identity` attribute to be set on the
4+
resulting ufunc, meaning it can be used for empty and multi-dimensional
5+
calls to :meth:`numpy.ufunc.reduce`.

‎numpy/core/_add_newdocs.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4180,7 +4180,7 @@
41804180

41814181
add_newdoc('numpy.core.umath', 'frompyfunc',
41824182
"""
4183-
frompyfunc(func, nin, nout)
4183+
frompyfunc(func, nin, nout, *[, identity])
41844184
41854185
Takes an arbitrary Python function and returns a NumPy ufunc.
41864186
@@ -4195,6 +4195,13 @@
41954195
The number of input arguments.
41964196
nout : int
41974197
The number of objects returned by `func`.
4198+
identity : object, optional
4199+
The value to use for the `~numpy.ufunc.identity` attribute of the resulting
4200+
object. If specified, this is equivalent to setting the underlying
4201+
C ``identity`` field to ``PyUFunc_IdentityValue``.
4202+
If omitted, the identity is set to ``PyUFunc_None``. Note that this is
4203+
_not_ equivalent to setting the identity to ``None``, which implies the
4204+
operation is reorderable.
41984205
41994206
Returns
42004207
-------

‎numpy/core/src/umath/umathmodule.c

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -70,9 +70,7 @@ object_ufunc_loop_selector(PyUFuncObject *ufunc,
7070
}
7171

7272
PyObject *
73-
ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *NPY_UNUSED(kwds)) {
74-
/* Keywords are ignored for now */
75-
73+
ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *kwds) {
7674
PyObject *function, *pyname = NULL;
7775
int nin, nout, i, nargs;
7876
PyUFunc_PyFuncData *fdata;
@@ -81,14 +79,18 @@ ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *NPY_UNUS
8179
Py_ssize_t fname_len = -1;
8280
void * ptr, **data;
8381
int offset[2];
82+
PyObject *identity = NULL; /* note: not the same semantics as Py_None */
83+
static char *kwlist[] = {"", "nin", "nout", "identity", NULL};
8484

85-
if (!PyArg_ParseTuple(args, "Oii:frompyfunc", &function, &nin, &nout)) {
85+
if (!PyArg_ParseTupleAndKeywords(args, kwds, "Oii|$O:frompyfunc", kwlist,
86+
&function, &nin, &nout, &identity)) {
8687
return NULL;
8788
}
8889
if (!PyCallable_Check(function)) {
8990
PyErr_SetString(PyExc_TypeError, "function must be callable");
9091
return NULL;
9192
}
93+
9294
nargs = nin + nout;
9395

9496
pyname = PyObject_GetAttrString(function, "__name__");
@@ -146,10 +148,10 @@ ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *NPY_UNUS
146148
/* Do a better job someday */
147149
doc = "dynamic ufunc based on a python function";
148150

149-
self = (PyUFuncObject *)PyUFunc_FromFuncAndData(
151+
self = (PyUFuncObject *)PyUFunc_FromFuncAndDataAndSignatureAndIdentity(
150152
(PyUFuncGenericFunction *)pyfunc_functions, data,
151-
types, /* ntypes */ 1, nin, nout, PyUFunc_None,
152-
str, doc, /* unused */ 0);
153+
types, /* ntypes */ 1, nin, nout, identity ? PyUFunc_IdentityValue : PyUFunc_None,
154+
str, doc, /* unused */ 0, NULL, identity);
153155

154156
if (self == NULL) {
155157
PyArray_free(ptr);

‎numpy/core/tests/test_umath.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2862,6 +2862,32 @@ def __new__(subtype, shape):
28622862
a = simple((3, 4))
28632863
assert_equal(a+a, a)
28642864

2865+
2866+
class TestFrompyfunc(object):
2867+
2868+
def test_identity(self):
2869+
def mul(a, b):
2870+
return a * b
2871+
2872+
# with identity=value
2873+
mul_ufunc = np.frompyfunc(mul, nin=2, nout=1, identity=1)
2874+
assert_equal(mul_ufunc.reduce([2, 3, 4]), 24)
2875+
assert_equal(mul_ufunc.reduce(np.ones((2, 2)), axis=(0, 1)), 1)
2876+
assert_equal(mul_ufunc.reduce([]), 1)
2877+
2878+
# with identity=None (reorderable)
2879+
mul_ufunc = np.frompyfunc(mul, nin=2, nout=1, identity=None)
2880+
assert_equal(mul_ufunc.reduce([2, 3, 4]), 24)
2881+
assert_equal(mul_ufunc.reduce(np.ones((2, 2)), axis=(0, 1)), 1)
2882+
assert_raises(ValueError, lambda: mul_ufunc.reduce([]))
2883+
2884+
# with no identity (not reorderable)
2885+
mul_ufunc = np.frompyfunc(mul, nin=2, nout=1)
2886+
assert_equal(mul_ufunc.reduce([2, 3, 4]), 24)
2887+
assert_raises(ValueError, lambda: mul_ufunc.reduce(np.ones((2, 2)), axis=(0, 1)))
2888+
assert_raises(ValueError, lambda: mul_ufunc.reduce([]))
2889+
2890+
28652891
def _check_branch_cut(f, x0, dx, re_sign=1, im_sign=-1, sig_zero_ok=False,
28662892
dtype=complex):
28672893
"""

0 commit comments

Comments
 (0)
0