summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/core/src/umath/ufunc_object.c127
-rw-r--r--numpy/core/tests/test_ufunc.py46
2 files changed, 160 insertions, 13 deletions
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 631999bc1..1091cdf6e 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -566,11 +566,12 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
NPY_ORDER *out_order,
NPY_CASTING *out_casting,
PyObject **out_extobj,
- PyObject **out_typetup,
- int *out_subok,
- PyArrayObject **out_wheremask,
- PyObject **out_axes,
- int *out_keepdims)
+ PyObject **out_typetup, /* type: Tuple[np.dtype] */
+ int *out_subok, /* bool */
+ PyArrayObject **out_wheremask, /* PyArray of bool */
+ PyObject **out_axes, /* type: List[Tuple[T]] */
+ PyObject **out_axis, /* type: T */
+ int *out_keepdims) /* bool */
{
int i, nargs;
int nin = ufunc->nin;
@@ -588,6 +589,9 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
if (out_axes != NULL) {
*out_axes = NULL;
}
+ if (out_axis != NULL) {
+ *out_axis = NULL;
+ }
if (out_wheremask != NULL) {
*out_wheremask = NULL;
}
@@ -827,9 +831,23 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
case 'a':
/* possible axes argument for generalized ufunc */
if (out_axes != NULL && strcmp(str, "axes") == 0) {
+ if (out_axis != NULL && *out_axis != NULL) {
+ PyErr_SetString(PyExc_TypeError,
+ "cannot specify both 'axis' and 'axes'");
+ goto fail;
+ }
Py_INCREF(value);
*out_axes = value;
-
+ bad_arg = 0;
+ }
+ else if (out_axis != NULL && strcmp(str, "axis") == 0) {
+ if (out_axes != NULL && *out_axes != NULL) {
+ PyErr_SetString(PyExc_TypeError,
+ "cannot specify both 'axis' and 'axes'");
+ goto fail;
+ }
+ Py_INCREF(value);
+ *out_axis = value;
bad_arg = 0;
}
break;
@@ -1045,6 +1063,10 @@ fail:
Py_XDECREF(*out_axes);
*out_axes = NULL;
}
+ if (out_axis != NULL) {
+ Py_XDECREF(*out_axis);
+ *out_axis = NULL;
+ }
return -1;
}
@@ -1891,6 +1913,27 @@ _has_output_coredims(PyUFuncObject *ufunc) {
}
/*
+ * Check whether the gufunc can be used with axis, i.e., that there is only
+ * a single, shared core dimension (which means that operands either have
+ * that dimension, or have no core dimensions). Returns 0 if all is fine,
+ * and sets an error and returns -1 if not.
+ */
+static int
+_check_axis_support(PyUFuncObject *ufunc) {
+ if (ufunc->core_num_dim_ix != 1) {
+ PyErr_Format(PyExc_TypeError,
+ "%s: axis can only be used with a single shared core "
+ "dimension, not with the %d distinct ones implied by "
+ "signature %s.",
+ ufunc_get_name_cstr(ufunc),
+ ufunc->core_num_dim_ix,
+ ufunc->core_signature);
+ return -1;
+ }
+ return 0;
+}
+
+/*
* Check whether the gufunc can be used with keepdims, i.e., that all its
* input arguments have the same number of core dimension, and all output
* arguments have no core dimensions. Returns 0 if all is fine, and sets
@@ -1905,7 +1948,7 @@ _check_keepdims_support(PyUFuncObject *ufunc) {
if (ufunc->core_num_dims[i] != (i < nin ? input_core_dims : 0)) {
PyErr_Format(PyExc_TypeError,
"%s does not support keepdims: its signature %s requires "
- "that %s %d has %d core dimensions, but keepdims can only "
+ "%s %d to have %d core dimensions, but keepdims can only "
"be used when all inputs have the same number of core "
"dimensions and all outputs have no core dimensions.",
ufunc_get_name_cstr(ufunc),
@@ -1920,6 +1963,42 @@ _check_keepdims_support(PyUFuncObject *ufunc) {
}
/*
+ * Translate axis to axes list of the form [(axis,), ...], with an
+ * empty tuple for operands without core dimensions.
+ * Returns an axes tuple or NULL on failure.
+ */
+static PyObject*
+_build_axes_tuple_from_axis(PyObject *axis, int core_num_dims[], int nop) {
+ int i;
+ PyObject *axes = NULL, *axis_tuple = NULL, *tuple;
+ PyObject *empty_tuple = PyTuple_New(0); /* cannot realistically fail */
+
+ axes = PyList_New(nop);
+ if (axes == NULL) {
+ return NULL;
+ }
+ axis_tuple = PyTuple_Pack(1, axis);
+ if (axis_tuple == NULL) {
+ goto fail;
+ }
+
+ for (i = 0; i < nop; i++) {
+ tuple = core_num_dims[i] == 1 ? axis_tuple : empty_tuple;
+ Py_INCREF(tuple);
+ PyList_SET_ITEM(axes, i, tuple);
+ }
+ Py_DECREF(axis_tuple);
+ Py_DECREF(empty_tuple);
+ return axes;
+
+fail:
+ Py_XDECREF(axis_tuple);
+ Py_XDECREF(axes);
+ Py_DECREF(empty_tuple);
+ return NULL;
+}
+
+/*
* Interpret a possible axes keyword argument, using it to fill the remap_axis
* array which maps default to actual axes for each operand, indexed as
* as remap_axis[iop][iaxis]. The default axis order has first all broadcast
@@ -2245,7 +2324,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
NPY_ORDER order = NPY_KEEPORDER;
/* Use the default assignment casting rule */
NPY_CASTING casting = NPY_DEFAULT_ASSIGN_CASTING;
- PyObject *extobj = NULL, *type_tup = NULL, *axes = NULL;
+ /* other possible keyword arguments */
+ PyObject *extobj = NULL, *type_tup = NULL, *axes = NULL, *axis = NULL;
int keepdims = -1;
if (ufunc == NULL) {
@@ -2270,10 +2350,12 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
NPY_UF_DBG_PRINT("Getting arguments\n");
- /* Get all the arguments */
+ /*
+ * Get all the arguments.
+ */
retval = get_ufunc_arguments(ufunc, args, kwds,
op, &order, &casting, &extobj,
- &type_tup, &subok, NULL, &axes, &keepdims);
+ &type_tup, &subok, NULL, &axes, &axis, &keepdims);
if (retval < 0) {
goto fail;
}
@@ -2288,6 +2370,12 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
goto fail;
}
}
+ if (axis != NULL) {
+ retval = _check_axis_support(ufunc);
+ if (retval < 0) {
+ goto fail;
+ }
+ }
/*
* If keepdims is set and true, signal all dimensions will be the same.
*/
@@ -2355,6 +2443,19 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
goto fail;
}
+ /*
+ * Translate axis to axes list of the form [(axis,), ...], with an
+ * empty tuple for operands without core dimensions.
+ */
+ if (axis) {
+ assert(axes == NULL); /* parser prevents passing in both axis & axes */
+ axes = _build_axes_tuple_from_axis(axis, core_num_dims, nop);
+ if (axes == NULL) {
+ retval = -1;
+ goto fail;
+ }
+ }
+
/* Possibly remap axes. */
if (axes) {
remap_axis = PyArray_malloc(sizeof(remap_axis[0]) * nop);
@@ -2714,8 +2815,10 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
Py_XDECREF(type_tup);
Py_XDECREF(extobj);
Py_XDECREF(axes);
+ Py_XDECREF(axis);
Py_XDECREF(full_args.in);
Py_XDECREF(full_args.out);
+ Py_XDECREF(axes);
NPY_UF_DBG_PRINT("Returning Success\n");
@@ -2735,8 +2838,10 @@ fail:
Py_XDECREF(type_tup);
Py_XDECREF(extobj);
Py_XDECREF(axes);
+ Py_XDECREF(axis);
Py_XDECREF(full_args.in);
Py_XDECREF(full_args.out);
+ Py_XDECREF(axes);
PyArray_free(remap_axis_memory);
PyArray_free(remap_axis);
return retval;
@@ -2812,7 +2917,7 @@ PyUFunc_GenericFunction(PyUFuncObject *ufunc,
/* Get all the arguments */
retval = get_ufunc_arguments(ufunc, args, kwds,
op, &order, &casting, &extobj,
- &type_tup, &subok, &wheremask, NULL, NULL);
+ &type_tup, &subok, &wheremask, NULL, NULL, NULL);
if (retval < 0) {
goto fail;
}
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index ac4346461..80237756d 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -726,6 +726,38 @@ class TestUfunc(object):
# should be able to deal with bad unrelated kwargs.
assert_raises(TypeError, mm, z, z, axes=[0, 1], parrot=True)
+ def test_axis_argument(self):
+ # inner1d signature: '(i),(i)->()'
+ inner1d = umt.inner1d
+ a = np.arange(27.).reshape((3, 3, 3))
+ b = np.arange(10., 19.).reshape((3, 1, 3))
+ c = inner1d(a, b)
+ assert_array_equal(c, (a * b).sum(-1))
+ c = inner1d(a, b, axis=-1)
+ assert_array_equal(c, (a * b).sum(-1))
+ out = np.zeros_like(c)
+ d = inner1d(a, b, axis=-1, out=out)
+ assert_(d is out)
+ assert_array_equal(d, c)
+ c = inner1d(a, b, axis=0)
+ assert_array_equal(c, (a * b).sum(0))
+ # Sanity check on innerwt.
+ a = np.arange(6).reshape((2, 3))
+ b = np.arange(10, 16).reshape((2, 3))
+ w = np.arange(20, 26).reshape((2, 3))
+ assert_array_equal(umt.innerwt(a, b, w, axis=0),
+ np.sum(a * b * w, axis=0))
+ # Check errors.
+ # Cannot pass in both axis and axes.
+ assert_raises(TypeError, inner1d, a, b, axis=0, axes=[0, 0])
+ # Not an integer.
+ assert_raises(TypeError, inner1d, a, b, axis=[0])
+ # more than 1 core dimensions.
+ mm = umt.matrix_multiply
+ assert_raises(TypeError, mm, a, b, axis=1)
+ # Regular ufuncs should not accept axis.
+ assert_raises(TypeError, np.add, 1., 1., axis=0)
+
def test_keepdims_argument(self):
# inner1d signature: '(i),(i)->()'
inner1d = umt.inner1d
@@ -741,7 +773,15 @@ class TestUfunc(object):
d = inner1d(a, b, keepdims=True, out=out)
assert_(d is out)
assert_array_equal(d, c)
- # Now combined with axes.
+ # Now combined with axis and axes.
+ c = inner1d(a, b, axis=-1, keepdims=False)
+ assert_array_equal(c, (a * b).sum(-1, keepdims=False))
+ c = inner1d(a, b, axis=-1, keepdims=True)
+ assert_array_equal(c, (a * b).sum(-1, keepdims=True))
+ c = inner1d(a, b, axis=0, keepdims=False)
+ assert_array_equal(c, (a * b).sum(0, keepdims=False))
+ c = inner1d(a, b, axis=0, keepdims=True)
+ assert_array_equal(c, (a * b).sum(0, keepdims=True))
c = inner1d(a, b, axes=[(-1,), (-1,), ()], keepdims=False)
assert_array_equal(c, (a * b).sum(-1))
c = inner1d(a, b, axes=[(-1,), (-1,), (-1,)], keepdims=True)
@@ -785,10 +825,12 @@ class TestUfunc(object):
w = np.arange(20, 26).reshape((2, 3))
assert_array_equal(umt.innerwt(a, b, w, keepdims=True),
np.sum(a * b * w, axis=-1, keepdims=True))
+ assert_array_equal(umt.innerwt(a, b, w, axis=0, keepdims=True),
+ np.sum(a * b * w, axis=0, keepdims=True))
# Check errors.
# Not a boolean
assert_raises(TypeError, inner1d, a, b, keepdims='true')
- # 1 core dimension only.
+ # More than 1 core dimension, and core output dimensions.
mm = umt.matrix_multiply
assert_raises(TypeError, mm, a, b, keepdims=True)
assert_raises(TypeError, mm, a, b, keepdims=False)
View Lorry Controller Status