8000 Merge pull request #5621 from jaimefrio/ufunc_out_tuple · numpy/numpy@bc034dc · GitHub
[go: up one dir, main page]
Skip to content

Commit bc034dc

Browse files
charrischarris
committed
Merge pull request #5621 from jaimefrio/ufunc_out_tuple
ENH: ufuncs can take a tuple of arrays as 'out' kwarg
2 parents a0ea059 + bb3ae05 commit bc034dc

File tree

4 files changed

+311
-100
lines changed

4 files changed

+311
-100
lines changed

‎doc/release/1.10.0-notes.rst

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,15 @@ Notably, this affect recarrays containing strings with whitespace, as trailing
7171
whitespace is trimmed from chararrays but kept in ndarrays of string type.
7272
Also, the dtype.type of nested structured fields is now inherited.
7373

74+
'out' keyword argument of ufuncs now accepts tuples of arrays
75+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
76+
When using the 'out' keyword argument of a ufunc, a tuple of arrays, one per
77+
ufunc output, can be provided. For ufuncs with a single output a single array
78+
is also a valid 'out' keyword argument. Previously a single array could be
79+
provided in the 'out' keyword argument, and it would be used as the first
80+
output for ufuncs with multiple outputs, is deprecated, and will result in a
81+
`DeprecationWarning` now and an error in the future.
82+
7483
New Features
7584
============
7685

‎doc/source/reference/ufuncs.rst

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -299,7 +299,20 @@ advanced usage and will not typically be used.
299299

300300
.. versionadded:: 1.6
301301

302-
The first output can provided as either a positional or a keyword parameter.
302+
The first output can be provided as either a positional or a keyword
303+
parameter. Keyword 'out' arguments are incompatible with positional
304+
ones.
305+
306+
..versionadded:: 1.10
307+
308+
The 'out' keyword argument is expected to be a tuple with one entry per
309+
output (which can be `None` for arrays to be allocated by the ufunc).
310+
For ufuncs with a single output, passing a single array (instead of a
311+
tuple holding a single array) is also valid.
312+
313+
Passing a single array in the 'out' keyword argument to a ufunc with
314+
multiple outputs is deprecated, and will raise a warning in numpy 1.10,
315+
and an error in a future release.
303316

304317
*where*
305318

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

Lines changed: 163 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -750,6 +750,35 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
750750
return -1;
751751
}
752752

753+
/*
754+
* Checks if 'obj' is a valid output array for a ufunc, i.e. it is
755+
* either None or a writeable array, increments its reference count
756+
* and stores a pointer to it in 'store'. Returns 0 on success, sets
757+
* an exception and returns -1 on failure.
758+
*/
759+
static i A3E2 nt
760+
_set_out_array(PyObject *obj, PyArrayObject **store)
761+
{
762+
if (obj == Py_None) {
763+
/* Translate None to NULL */
764+
return 0;
765+
}
766+
if PyArray_Check(obj) {
767+
/* If it's an array, store it */
768+
if (PyArray_FailUnlessWriteable((PyArrayObject *)obj,
769+
"output array") < 0) {
770+
return -1;
771+
}
772+
Py_INCREF(obj);
773+
*store = (PyArrayObject *)obj;
774+
775+
return 0;
776+
}
777+
PyErr_SetString(PyExc_TypeError, "return arrays must be of ArrayType");
778+
779+
return -1;
780+
}
781+
753782
/********* GENERIC UFUNC USING ITERATOR *********/
754783

755784
/*
@@ -759,17 +788,20 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
759788
* non-zero references in out_op. This
760789
* function does not do its own clean-up.
761790
*/
762-
static int get_ufunc_arguments(PyUFuncObject *ufunc,
763-
PyObject *args, PyObject *kwds,
764-
PyArrayObject **out_op,
765-
NPY_ORDER *out_order,
766-
NPY_CASTING *out_casting,
767-
PyObject **out_extobj,
768-
PyObject **out_typetup,
769-
int *out_subok,
770-
PyArrayObject **out_wheremask)
791+
static int
792+
get_ufunc_arguments(PyUFuncObject *ufunc,
793+
PyObject *args, PyObject *kwds,
794+
PyArrayObject **out_op,
795+
NPY_ORDER *out_order,
796+
NPY_CASTING *out_casting,
797+
PyObject **out_extobj,
798+
PyObject **out_typetup,
799+
int *out_subok,
800+
PyArrayObject **out_wheremask)
771801
{
772-
int i, nargs, nin = ufunc->nin;
802+
int i, nargs;
803+
int nin = ufunc->nin;
804+
int nout = ufunc->nout;
773805
PyObject *obj, *context;
774806
PyObject *str_key_obj = NULL;
775807
const char *ufunc_name;
@@ -878,23 +910,7 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
878910
/* Get positional output arguments */
879911
for (i = nin; i < nargs; ++i) {
880912
obj = PyTuple_GET_ITEM(args, i);
881-
/* Translate None to NULL */
882-
if (obj == Py_None) {
883-
continue;
884-
}
885-
/* If it's an array, can use it */
886-
if (PyArray_Check(obj)) {
887-
if (PyArray_FailUnlessWriteable((PyArrayObject *)obj,
888-
"output array") < 0) {
889-
return -1;
890-
}
891-
Py_INCREF(obj);
892-
out_op[i] = (PyArrayObject *)obj;
893-
}
894-
else {
895-
PyErr_SetString(PyExc_TypeError,
896-
"return arrays must be "
897-
"of ArrayType");
913+
if (_set_out_array(obj, out_op + i) < 0) {
898914
return -1;
899915
}
900916
}
@@ -929,7 +945,7 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
929945
switch (str[0]) {
930946
case 'c':
931947
/* Provides a policy for allowed casting */
932-
if (strncmp(str,"casting",7) == 0) {
948+
if (strncmp(str, "casting", 7) == 0) {
933949
if (!PyArray_CastingConverter(value, out_casting)) {
934950
goto fail;
935951
}
@@ -938,7 +954,7 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
938954
break;
939955
case 'd':
940956
/* Another way to specify 'sig' */
941-
if (strncmp(str,"dtype",5) == 0) {
957+
if (strncmp(str, "dtype", 5) == 0) {
942958
/* Allow this parameter to be None */
943959
PyArray_Descr *dtype;
944960
if (!PyArray_DescrConverter2(value, &dtype)) {
@@ -960,35 +976,74 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
960976
* Overrides the global parameters buffer size,
961977
* error mask, and error object
962978
*/
963-
if (strncmp(str,"extobj",6) == 0) {
979+
if (strncmp(str, "extobj", 6) == 0) {
964980
*out_extobj = value;
965981
bad_arg = 0;
966982
}
967983
break;
968984
case 'o':
969-
/* First output may be specified as a keyword parameter */
970-
if (strncmp(str,"out",3) == 0) {
971-
if (out_op[nin] != NULL) {
985+
/*
986+
* Output arrays may be specified as a keyword argument,
987+
* either as a single array or None for single output
988+
* ufuncs, or as a tuple of arrays and Nones.
989+
*/
990+
if (strncmp(str, "out", 3) == 0) {
991+
if (nargs > nin) {
972992
PyErr_SetString(PyExc_ValueError,
973993
"cannot specify 'out' as both a "
974994
"positional and keyword argument");
975995
goto fail;
976996
}
977-
978-
if (PyArray_Check(value)) {
979-
const char *name = "output array";
980-
PyArrayObject *value_arr = (PyArrayObject *)value;
981-
if (PyArray_FailUnlessWriteable(value_arr, name) < 0) {
997+
if (PyTuple_Check(value)) {
998+
if (PyTuple_GET_SIZE(value) != nout) {
999+
PyErr_SetString(PyExc_ValueError,
1000+
"The 'out' tuple must have exactly "
1001+
"one entry per ufunc output");
1002+
goto fail;
1003+
}
1004+
/* 'out' must be a tuple of arrays and Nones */
1005+
for(i = 0; i < nout; ++i) {
1006+
PyObject *val = PyTuple_GET_ITEM(value, i);
1007+
if (_set_out_array(val, out_op+nin+i) < 0) {
1008+
goto fail;
1009+
}
1010+
}
1011+
}
1012+
else if (nout == 1) {
1013+
/* Can be an array if it only has one output */
1014+
if (_set_out_array(value, out_op + nin) < 0) {
9821015
goto fail;
9831016
}
984-
Py_INCREF(value);
985-
out_op[nin] = (PyArrayObject *)value;
9861017
}
9871018
else {
988-
PyErr_SetString(PyExc_TypeError,
989-
"return arrays must be "
990-
"of ArrayType");
991-
goto fail;
1019+
/*
1020+
* If the deprecated behavior is ever removed,
1021+
* keep only the else branch of this if-else
1022+
*/
1023+
if (PyArray_Check(value) || value == Py_None) {
1024+
if (DEPRECATE("passing a single array to the "
1025+
"'out' keyword argument of a "
1026+
"ufunc with\n"
1027+
"more than one output will "
1028+
"result in an error in the "
1029+
"future") < 0) {
1030+
/* The future error message */
1031+
PyErr_SetString(PyExc_TypeError,
1032+
"'out' must be a tuple of arrays");
1033+
goto fail;
1034+
}
1035+
if (_set_out_array(value, out_op+nin) < 0) {
1036+
goto fail;
1037+
}
1038+
}
1039+
else {
1040+
PyErr_SetString(PyExc_TypeError,
1041+
nout > 1 ? "'out' must be a tuple "
1042+
"of arrays" :
1043+
"'out' must be an array or a "
1044+
"tuple of a single array");
1045+
goto fail;
1046+
}
9921047
}
9931048
bad_arg = 0;
9941049
}
@@ -3945,6 +4000,38 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args,
39454000
return PyArray_Return(ret);
39464001
}
39474002

4003+
/*
4004+
* Returns an incref'ed pointer to the proper wrapping object for a
4005+
* ufunc output argument, given the output argument 'out', and the
4006+
* input's wrapping function, 'wrap'.
4007+
*/
4008+
static PyObject*
4009+
_get_out_wrap(PyObject *out, PyObject *wrap) {
4010+
PyObject *owrap;
4011+
4012+
if (out == Py_No 10000 ne) {
4013+
/* Iterator allocated outputs get the input's wrapping */
4014+
Py_XINCREF(wrap);
4015+
return wrap;
4016+
}
4017+
if (PyArray_CheckExact(out)) {
4018+
/* None signals to not call any wrapping */
4019+
Py_RETURN_NONE;
4020+
}
4021+
/*
4022+
* For array subclasses use their __array_wrap__ method, or the
4023+
* input's wrapping if not available
4024+
*/
4025+
owrap = PyObject_GetAttr(out, npy_um_str_array_wrap);
4026+
if (owrap == NULL || !PyCallable_Check(owrap)) {
4027+
Py_XDECREF(owrap);
4028+
owrap = wrap;
4029+
Py_XINCREF(wrap);
4030+
PyErr_Clear();
4031+
}
4032+
return owrap;
4033+
}
4034+
39484035
/*
39494036
* This function analyzes the input arguments
39504037
* and determines an appropriate __array_wrap__ function to call
@@ -3966,7 +4053,7 @@ _find_array_wrap(PyObject *args, PyObject *kwds,
39664053
PyObject **output_wrap, int nin, int nout)
39674054
{
39684055
Py_ssize_t nargs;
3969-
int i;
4056+
int i, idx_offset, start_idx;
39704057
int np = 0;
39714058
PyObject *with_wrap[NPY_MAXARGS], *wraps[NPY_MAXARGS];
39724059
PyObject *obj, *wrap = NULL;
@@ -4043,45 +4130,45 @@ _find_array_wrap(PyObject *args, PyObject *kwds,
40434130
*/
40444131
handle_out:
40454132
nargs = PyTuple_GET_SIZE(args);
4046-
for (i = 0; i < nout; i++) {
4047-
int j = nin + i;
4048-
int incref = 1;
4049-
output_wrap[i] = wrap;
4050-
obj = NULL;
4051-
if (j < nargs) {
4052-
obj = PyTuple_GET_ITEM(args, j);
4053-
/* Output argument one may also be in a keyword argument */
4054-
if (i == 0 && obj == Py_None && kwds != NULL) {
4055-
obj = PyDict_GetItem(kwds, npy_um_str_out);
4056-
}
4133+
/* Default is using positional arguments */
4134+
obj = args;
4135+
idx_offset = nin;
4136+
start_idx = 0;
4137+
if (nin == nargs && kwds != NULL) {
4138+
/* There may be a keyword argument we can use instead */
4139+
obj = PyDict_GetItem(kwds, npy_um_str_out);
4140+
if (obj == NULL) {
4141+
/* No, go back to positional (even though there aren't any) */
4142+
obj = args;
40574143
}
4058-
/* Output argument one may also be in a keyword argument */
4059-
else if (i == 0 && kwds != NULL) {
4060-
obj = PyDict_GetItem(kwds, npy_um_str_out);
4061-
}
4062-
4063-
if (obj != Py_None && obj != NULL) {
4064-
if (PyArray_CheckExact(obj)) {
4065-
/* None signals to not call any wrapping */
4066-
output_wrap[i] = Py_None;
4144+
else {
4145+
idx_offset = 0;
4146+
if (PyTuple_Check(obj)) {
4147+
/* If a tuple, must have all nout items */
4148+
nargs = nout;
40674149
}
40684150
else {
4069-
PyObject *owrap = PyObject_GetAttr(obj, npy_um_str_array_wrap);
4070-
incref = 0;
4071-
if (!(owrap) || !(PyCallable_Check(owrap))) {
4072-
Py_XDECREF(owrap);
4073-
owrap = wrap;
4074-
incref = 1;
4075-
PyErr_Clear();
4076-
}
4077-
output_wrap[i] = owrap;
4151+
/* If the kwarg is not a tuple then it is an array (or None) */
4152+
output_wrap[0] = _get_out_wrap(obj, wrap);
4153+
start_idx = 1;
4154+
nargs = 1;
40784155
}
40794156
}
4157+
}
40804158

4081-
if (incref) {
4082-
Py_XINCREF(output_wrap[i]);
4159+
for (i = start_idx; i < nout; ++i) {
4160+
int j = idx_offset + i;
4161+
4162+
if (j < nargs) {
4163+
output_wrap[i] = _get_out_wrap(PyTuple_GET_ITEM(obj, j),
4164+
wrap);
4165+
}
4166+
else {
4167+
output_wrap[i] = wrap;
4168+
Py_XINCREF(wrap);
40834169
}
40844170
}
4171+
40854172
Py_XDECREF(wrap);
40864173
return;
40874174
}

0 commit comments

Comments
 (0)
0