8000 Merge pull request #11018 from mhvk/gufunc-axis-and-keepdims · numpy/numpy@57d17c3 · GitHub
[go: up one dir, main page]
Skip to content

Commit 57d17c3

Browse files
charrischarris
authored
Merge pull request #11018 from mhvk/gufunc-axis-and-keepdims
ENH: Implement axis for generalized ufuncs.
2 parents 5cbb982 + b02fa32 commit 57d17c3

File tree

7 files changed

+271
-24
lines changed

7 files changed

+271
-24
lines changed

‎doc/release/1.15.0-notes.rst

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ combining these 5 compiled builds products into a single "fat" binary.
193193
``return_indices`` keyword added for ``np.intersect1d``
194194
-------------------------------------------------------
195195
New keyword ``return_indices`` returns the indices of the two input arrays
196-
that correspond to the common elements.
196+
that correspond to the common elements.
197197

198198
``np.quantile`` and ``np.nanquantile``
199199
--------------------------------------
@@ -359,8 +359,8 @@ Increased performance in ``random.permutation`` for multidimensional arrays
359359
``permutation`` uses the fast path in ``random.shuffle`` for all input
360360
array dimensions. Previously the fast path was only used for 1-d arrays.
361361

362-
Generalized ufuncs now accept ``axes`` and ``keepdims`` arguments
363-
-----------------------------------------------------------------
362+
Generalized ufuncs now accept ``axes``, ``axis`` and ``keepdims`` arguments
363+
---------------------------------------------------------------------------
364364
One can control over which axes a generalized ufunc operates by passing in an
365365
``axes`` argument, a list of tuples with indices of particular axes. For
366366
instance, for a signature of ``(i,j),(j,k)->(i,k)`` appropriate for matrix
@@ -376,12 +376,19 @@ tuples can be omitted. Hence, for a signature of ``(i),(i)->()`` appropriate
376376
for an inner product, one could pass in ``axes=[0, 0]`` to indicate that the
377377
vectors are stored in the first dimensions of the two inputs arguments.
378378

379+
As a short-cut for generalized ufuncs that are similar to reductions, i.e.,
380+
that act on a single, shared core dimension such as the inner product example
381+
above, one can pass an ``axis`` argument. This is equivalent to passing in
382+
``axes`` with identical entries for all arguments with that core dimension
383+
(e.g., for the example above, ``axes=[(axis,), (axis,)]``).
384+
379385
Furthermore, like for reductions, for generalized ufuncs that have inputs that
380386
all have the same number of core dimensions and outputs with no core dimension,
381387
one can pass in ``keepdims`` to leave a dimension with size 1 in the outputs,
382388
thus allowing proper broadcasting against the original inputs. The location of
383389
the extra dimension can be controlled with ``axes``. For instance, for the
384390
inner-product example, ``keepdims=True, axes=[-2, -2, -2]`` would act on the
391+
inner-product example, ``keepdims=True, axis=-2`` would act on the
385392
one-but-last dimension of the input arguments, and leave a size 1 dimension in
386393
that place in the output.
387394

@@ -411,6 +418,9 @@ is the same as::
411418
``np.put_along_axis`` acts as the dual operation for writing to these indices
412419
within an array.
413420

421+
.. note:: Implementations of ``__array_ufunc__`` should ensure that they can
422+
handle either ``axis`` or ``axes``. In future, we may convert
423+
``axis`` to ``axes`` before passing it on.
414424

415425
Changes
416426
=======

‎doc/source/reference/ufuncs.rst

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -360,6 +360,17 @@ advanced usage and will not typically be used.
360360
and for generalized ufuncs for which all outputs are scalars, the output
361361
tuples can be omitted.
362362

363+
*axis*
364+
365+
.. versionadded:: 1.15
366+
367+
A single axis over which a generalized ufunc should operate. This is a
368+
short-cut for ufuncs that operate over a single, shared core dimension,
369+
equivalent to passing in ``axes`` with entries of ``(axis,)`` for each
370+
single-core-dimension argument and ``()`` for all others. For instance,
371+
for a signature ``(i),(i)->()``, it is equivalent to passing in
372+
``axes=[(axis,), (axis,), ()]``.
373+
363374
*keepdims*
364375

365376
.. versionadded:: 1.15
@@ -370,7 +381,7 @@ advanced usage and will not typically be used.
370381
ufuncs that operate on inputs that all have the same number of core
371382
dimensions and with outputs that have no core dimensions , i.e., with
372383
signatures like ``(i),(i)->()`` or ``(m,m)->()``. If used, the location of
373-
the dimensions in the output can be controlled with ``axes``.
384+
the dimensions in the output can be controlled with ``axes`` and ``axis``.
374385

375386
*casting*
376387

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,38 @@ static void
253253

254254
/**end repeat**/
255255

256+
char *cumsum_signature = "(i)->(i)";
257+
258+
/*
259+
* This implements the function
260+
* out[n] = sum_i^n in[i]
261+
*/
262+
263+
/**begin repeat
264+
265+
#TYPE=LONG,DOUBLE#
266+
#typ=npy_long,npy_double#
267+
*/
268+
269+
static void
270+
@TYPE@_cumsum(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
271+
{
272+
INIT_OUTER_LOOP_2
273+
npy_intp di = dimensions[0];
274+
npy_intp i;
275+
npy_intp is=steps[0], os=steps[1];
276+
BEGIN_OUTER_LOOP_2
277+
char *ip=args[0], *op=args[1];
278+
@typ@ cumsum = 0;
279+
for (i = 0; i < di; i++, ip += is, op += os) {
280+
cumsum += (*(@typ@ *)ip);
281+
*(@typ@ *)op = cumsum;
282+
}
283+
END_OUTER_LOOP
284+
}
285+
286+
/**end repeat**/
287+
256288

257289
static PyUFuncGenericFunction inner1d_functions[] = { LONG_inner1d, DOUBLE_inner1d };
258290
static void * inner1d_data[] = { (void *)NULL, (void *)NULL };
@@ -270,6 +302,10 @@ static void *eucldiean_pdist_data[] = { (void *)NULL, (void *)NULL };
270302
static char euclidean_pdist_signatures[] = { NPY_FLOAT, NPY_FLOAT,
271303
NPY_DOUBLE, NPY_DOUBLE };
272304

305+
static PyUFuncGenericFunction cumsum_functions[] = { LONG_cumsum, DOUBLE_cumsum };
306+
static void * cumsum_data[] = { (void *)NULL, (void *)NULL };
307+
static char cumsum_signatures[] = { NPY_LONG, NPY_LONG, NPY_DOUBLE, NPY_DOUBLE };
308+
273309

274310
static int
275311
addUfuncs(PyObject *dictionary) {
@@ -321,6 +357,16 @@ addUfuncs(PyObject *dictionary) {
321357
}
322358
PyDict_SetItemString(dictionary, "euclidean_pdist", f);
323359
Py_DECREF(f);
360+
f = PyUFunc_FromFuncAndDataAndSignature(cumsum_functions,
361+
cumsum_data, cumsum_signatures,
362+
2, 1, 1, PyUFunc_None, "cumsum",
363+
"Cumulative sum of the input (n)->(n)\n",
364+
0, cumsum_signature);
365+
if (f == NULL) {
366+
return -1;
367+
}
368+
PyDict_SetItemString(dictionary, "cumsum", f);
369+
Py_DECREF(f);
324370
f = PyUFunc_FromFuncAndDataAndSignature(inner1d_functions, inner1d_data,
325371
inner1d_signatures, 2, 2, 1, PyUFunc_None, "inner1d_no_doc",
326372
NULL,

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ normalize___call___args(PyUFuncObject *ufunc, PyObject *args,
5151
npy_intp nin = ufunc->nin;
5252
npy_intp nout = ufunc->nout;
5353
npy_intp nargs = PyTuple_GET_SIZE(args);
54+
npy_intp nkwds = PyDict_Size(*normal_kwds);
5455
PyObject *obj;
5556

5657
if (nargs < nin) {
@@ -74,7 +75,7 @@ normalize___call___args(PyUFuncObject *ufunc, PyObject *args,
7475

7576
/* If we have more args than nin, they must be the output variables.*/
7677
if (nargs > nin) {
77-
if(PyDict_GetItemString(*normal_kwds, "out")) {
78+
if(nkwds > 0 && PyDict_GetItemString(*normal_kwds, "out")) {
7879
PyErr_Format(PyExc_TypeError,
7980
"argument given by name ('out') and position "
8081
"(%"NPY_INTP_FMT")", nin);
@@ -112,8 +113,15 @@ normalize___call___args(PyUFuncObject *ufunc, PyObject *args,
112113
Py_DECREF(obj);
113114
}
114115
}
116+
/* gufuncs accept either 'axes' or 'axis', but not both */
117+
if (nkwds >= 2 && (PyDict_GetItemString(*normal_kwds, "axis") &&
118+
PyDict_GetItemString(*normal_kwds, "axes"))) {
119+
PyErr_SetString(PyExc_TypeError,
120+
"cannot specify both 'axis' and 'axes'");
121+
return -1;
122+
}
115123
/* finally, ufuncs accept 'sig' or 'signature' normalize to 'signature' */
116-
return normalize_signature_keyword(*normal_kwds);
124+
return nkwds == 0 ? 0 : normalize_signature_keyword(*normal_kwds);
117125
}
118126

119127
static int

0 commit comments

Comments
 (0)
0