8000 Merge pull request #2953 from ContinuumIO/gufunc-fix · numpy/numpy@985b267 · GitHub
[go: up one dir, main page]
Skip to content

Commit 985b267

Browse files
charrischarris
committed
Merge pull request #2953 from ContinuumIO/gufunc-fix
generalized ufunc signature problem fix
2 parents d21281a + f2c81b4 commit 985b267

File tree

2 files changed

+199
-50
lines changed

2 files changed

+199
-50
lines changed

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

Lines changed: 190 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,12 @@
7171
static int
7272
_does_loop_use_arrays(void *data);
7373

74+
static int
75+
assign_reduce_identity_zero(PyArrayObject *result);
76+
77+
static int
78+
assign_reduce_identity_one(PyArrayObject *result);
79+
7480
/*
7581
* fpstatus is the ufunc_formatted hardware status
7682
* errmask is the handling mask specified by the user.
@@ -1643,19 +1649,20 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
16431649
PyArrayObject **op)
16441650
{
16451651
int nin, nout;
1646-
int i, idim, nop;
1652+
int i, j, idim, nop;
16471653
char *ufunc_name;
16481654
int retval = -1, subok = 1;
16491655
int needs_api = 0;
16501656

16511657
PyArray_Descr *dtypes[NPY_MAXARGS];
16521658

16531659
/* Use remapped axes for generalized ufunc */
1654-
int broadcast_ndim, op_ndim;
1660+
int broadcast_ndim, iter_ndim;
16551661
int op_axes_arrays[NPY_MAXARGS][NPY_MAXDIMS];
16561662
int *op_axes[NPY_MAXARGS];
16571663

16581664
npy_uint32 op_flags[NPY_MAXARGS];
1665+
npy_intp iter_shape[NPY_MAXARGS];
16591666

16601667
NpyIter *iter = NULL;
16611668

@@ -1673,7 +1680,9 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
16731680
npy_intp *inner_strides = NULL;
16741681

16751682
npy_intp *inner_strides_tmp, *ax_strides_tmp[NPY_MAXDIMS];
1676-
int core_dim_ixs_size, *core_dim_ixs;
1683+
/* The sizes of the core dimensions (# entries is ufunc->core_num_dim_ix) */
1684+
npy_intp *core_dim_sizes = inner_dimensions + 1;
1685+
int core_dim_ixs_size;
16771686

16781687
/* The __array_prepare__ function to call for each output */
16791688
PyObject *arr_prep[NPY_MAXARGS];
@@ -1719,26 +1728,107 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
17191728
goto fail;
17201729
}
17211730

1722-
/* Figure out the number of dimensions needed by the iterator */
1731+
/*
1732+
* Figure out the number of iteration dimensions, which
1733+
* is the broadcast result of all the input non-core
1734+
* dimensions.
1735+
*/
17231736
broadcast_ndim = 0;
17241737
for (i = 0; i < nin; ++i) {
17251738
int n = PyArray_NDIM(op[i]) - ufunc->core_num_dims[i];
17261739
if (n > broadcast_ndim) {
17271740
broadcast_ndim = n;
17281741
}
17291742
}
1730-
op_ndim = broadcast_ndim + ufunc->core_num_dim_ix;
1731-
if (op_ndim > NPY_MAXDIMS) {
1743+
1744+
/*
1745+
* Figure out the number of iterator creation dimensions,
1746+
* which is the broadcast dimensions + all the core dimensions of
1747+
* the outputs, so that the iterator can allocate those output
1748+
* dimensions following the rules of order='F', for example.
1749+
*/
1750+
iter_ndim = broadcast_ndim;
1751+
for (i = nin; i < nop; ++i) {
1752+
iter_ndim += ufunc->core_num_dims[i];
1753+
}
1754+
if (iter_ndim > NPY_MAXDIMS) {
17321755
PyErr_Format(PyExc_ValueError,
17331756
"too many dimensions for generalized ufunc %s",
17341757
ufunc_name);
17351758
retval = -1;
17361759
goto fail;
17371760
}
17381761

1762+
/*
1763+
* Validate the core dimensions of all the operands,
1764+
* and collect all of the labeled core dimension sizes
1765+
* into the array 'core_dim_sizes'. Initialize them to
1766+
* 1, for example in the case where the operand broadcasts
1767+
* to a core dimension, it won't be visited.
1768+
*/
1769+
for (i = 0; i < ufunc->core_num_dim_ix; ++i) {
1770+
core_dim_sizes[i] = 1;
1771+
}
1772+
for (i = 0; i < nop; ++i) {
1773+
if (op[i] != NULL) {
1774+
int dim_offset = ufunc->core_offsets[i];
1775+
int num_dims = ufunc->core_num_dims[i];
1776+
int core_start_dim = PyArray_NDIM(op[i]) - num_dims;
1777+
/* Make sure any output operand has enough dimensions */
1778+
if (i >= nin && core_start_dim < 0) {
1779+
PyErr_Format(PyExc_ValueError,
1780+
"%s: Output operand %d does not have enough dimensions "
1781+
"(has %d, gufunc core with signature %s "
1782+
"requires %d)",
1783+
ufunc_name, i - nin, PyArray_NDIM(op[i]),
1784+
ufunc->core_signature, num_dims);
1785+
retval = -1;
1786+
goto fail;
1787+
}
1788+
1789+
/*
1790+
* Make sure each core dimension matches all other core
1791+
* dimensions with the same label
1792+
*
1793+
* NOTE: For input operands, core_start_dim may be negative.
1794+
* In that case, the operand is being broadcast onto
1795+
* core dimensions. For example, a scalar will broadcast
1796+
* to fit any core signature.
1797+
*/
1798+
if (core_start_dim >= 0) {
1799+
idim = 0;
1800+
} else {
1801+
idim = -core_start_dim;
1802+
}
1803+
for (; idim < num_dims; ++idim) {
1804+
int core_dim_index = ufunc->core_dim_ixs[dim_offset + idim];
1805+
npy_intp op_dim_size =
1806+
PyArray_SHAPE(op[i])[core_start_dim + idim];
1807+
if (core_dim_sizes[core_dim_index] == 1) {
1808+
core_dim_sizes[core_dim_index] = op_dim_size;
1809+
} else if ((i >= nin || op_dim_size != 1) &&
1810+
core_dim_sizes[core_dim_index] != op_dim_size) {
1811+
PyErr_Format(PyExc_ValueError,
1812+
"%s: Operand %d has a mismatch in its core "
1813+
"dimension %d, with gufunc signature %s "
1814+
"(size %d is different from %d)",
1815+
ufunc_name, i, idim, ufunc->core_signature,
1816+
op_dim_size, core_dim_sizes[core_dim_index]);
1817+
retval = -1;
1818+
goto fail;
1819+
}
1820+
}
1821+
}
1822+
}
1823+
1824+
/* Fill in the initial part of 'iter_shape' */
1825+
for (idim = 0; idim < broadcast_ndim; ++idim) {
1826+
iter_shape[idim] = -1;
1827+
}
1828+
17391829
/* Fill in op_axes for all the operands */
1830+
j = broadcast_ndim;
17401831
core_dim_ixs_size = 0;
1741-
core_dim_ixs = ufunc->core_dim_ixs;
17421832
for (i = 0; i < nop; ++i) {
17431833
int n;
17441834
if (op[i]) {
@@ -1760,22 +1850,27 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
17601850
op_axes_arrays[i][idim] = -1;
17611851
}
17621852
}
1763-
/* Use the signature information for the rest */
1764-
for (idim = broadcast_ndim; idim < op_ndim; ++idim) {
1853+
1854+
/* Any output core dimensions shape should be ignored */
1855+
for (idim = broadcast_ndim; idim < iter_ndim; ++idim) {
17651856
op_axes_arrays[i][idim] = -1;
17661857
}
1767-
for (idim = 0; idim < ufunc->core_num_dims[i]; ++idim) {
1768-
if (n + idim >= 0) {
1769-
op_axes_arrays[i][broadcast_ndim + core_dim_ixs[idim]] =
1770-
n + idim;
1771-
}
1772-
else {
1773-
op_axes_arrays[i][broadcast_ndim + core_dim_ixs[idim]] = -1;
1858+
1859+
/* Except for when it belongs to this output */
1860+
if (i >= nin) {
1861+
int dim_offset = ufunc->core_offsets[i];
1862+
int num_dims = ufunc->core_num_dims[i];
1863+
/* Fill in 'iter_shape' and 'op_axes' for this output */
1864+
for (idim = 0; idim < num_dims; ++idim) {
1865+
iter_shape[j] = core_dim_sizes[
1866+
ufunc->core_dim_ixs[dim_offset + idim]];
1867+
op_axes_arrays[i][j] = n + idim;
1868+
++j;
17741869
}
17751870
}
1776-
core_dim_ixs_size += ufunc->core_num_dims[i];
1777-
core_dim_ixs += ufunc->core_num_dims[i];
1871+
17781872
op_axes[i] = op_axes_arrays[i];
1873+
core_dim_ixs_size += ufunc->core_num_dims[i];
17791874
}
17801875

17811876
/* Get the buffersize, errormask, and error object globals */
@@ -1881,12 +1976,26 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
18811976
NPY_ITER_NO_BROADCAST;
18821977
}
18831978

1979+
/*
1980+
* If there are no iteration dimensions, create a fake one
1981+
* so that the scalar edge case works right.
1982+
*/
1983+
if (iter_ndim == 0) {
1984+
iter_ndim = 1;
1985+
iter_shape[0] = 1;
1986+
for (i = 0; i < nop; ++i) {
1987+
op_axes[i][0] = -1;
1988+
}
1989+
}
1990+
18841991
/* Create the iterator */
18851992
iter = NpyIter_AdvancedNew(nop, op, NPY_ITER_MULTI_INDEX|
18861993
NPY_ITER_REFS_OK|
1887-
NPY_ITER_REDUCE_OK,
1994+
NPY_ITER_REDUCE_OK|
1995+
NPY_ITER_ZEROSIZE_OK,
18881996
order, NPY_UNSAFE_CASTING, op_flags,
1889-
dtypes, op_ndim, op_axes, NULL, 0);
1997+
dtypes, iter_ndim,
1998+
op_axes, iter_shape, 0);
18901999
if (iter == NULL) {
18912000
retval = -1;
18922001
goto fail;
@@ -1906,37 +2015,38 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
19062015
*/
19072016
inner_strides = (npy_intp *)PyArray_malloc(
19082017
NPY_SIZEOF_INTP * (nop+core_dim_ixs_size));
1909-
/* The strides after the first nop match core_dim_ixs */
1910-
core_dim_ixs = ufunc->core_dim_ixs;
1911-
inner_strides_tmp = inner_strides + nop;
1912-
for (idim = 0; idim < ufunc->core_num_dim_ix; ++idim) {
1913-
ax_strides_tmp[idim] = NpyIter_GetAxisStrideArray(iter,
1914-
broadcast_ndim+idim);
1915-
if (ax_strides_tmp[idim] == NULL) {
1916-
retval = -1;
1917-
goto fail;
1918-
}
1919-
}
2018+
/* Copy the strides after the first nop */
2019+
idim = nop;
19202020
for (i = 0; i < nop; ++i) {
1921-
for (idim = 0; idim < ufunc->core_num_dims[i]; ++idim) {
1922-
inner_strides_tmp[idim] = ax_strides_tmp[core_dim_ixs[idim]][i];
2021+
int dim_offset = ufunc->core_offsets[i];
2022+
int num_dims = ufunc->core_num_dims[i];
2023+
int core_start_dim = PyArray_NDIM(op[i]) - num_dims;
2024+
/*
2025+
* Need to use the arrays in the iterator, not op, because
2026+
* a copy with a different-sized type may have been made.
2027+
*/
2028+
PyArrayObject *arr = NpyIter_GetOperandArray(iter)[i];
2029+
npy_intp *shape = PyArray_SHAPE(arr);
2030+
npy_intp *strides = PyArray_STRIDES(arr);
2031+
for (j = 0; j < num_dims; ++j) {
2032+
if (core_start_dim + j >= 0) {
2033+
/*
2034+
* Force the stride to zero when the shape is 1, sot
2035+
* that the broadcasting works right.
2036+
*/
2037+
if (shape[core_start_dim + j] != 1) {
2038+
inner_strides[idim++] = strides[core_start_dim + j];
2039+
} else {
2040+
inner_strides[idim++] = 0;
2041+
}
2042+
} else {
2043+
inner_strides[idim++] = 0;
2044+
}
19232045
}
1924-
1925-
core_dim_ixs += ufunc->core_num_dims[i];
1926-
inner_strides_tmp += ufunc->core_num_dims[i];
1927-
}
1928-
1929-
/* Set up the inner dimensions array */
1930-
if (NpyIter_GetShape(iter, inner_dimensions) != NPY_SUCCEED) {
1931-
retval = -1;
1932-
goto fail;
19332046
}
1934-
/* Move the core dimensions to start at the second element */
1935-
memmove(&inner_dimensions[1], &inner_dimensions[broadcast_ndim],
1936-
NPY_SIZEOF_INTP * ufunc->core_num_dim_ix);
19372047

1938-
/* Remove all the core dimensions from the iterator */
1939-
for (i = 0; i < ufunc->core_num_dim_ix; ++i) {
2048+
/* Remove all the core output dimensions from the iterator */
2049+
for (i = broadcast_ndim; i < iter_ndim; ++i) {
19402050
if (NpyIter_RemoveAxis(iter, broadcast_ndim) != NPY_SUCCEED) {
19412051
retval = -1;
19422052
goto fail;
@@ -1971,8 +2081,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
19712081

19722082
NPY_UF_DBG_PRINT("Executing inner loop\n");
19732083

1974-
/* Do the ufunc loop */
19752084
if (NpyIter_GetIterSize(iter) != 0) {
2085+
/* Do the ufunc loop */
19762086
NpyIter_IterNextFunc *iternext;
19772087
char **dataptr;
19782088
npy_intp *count_ptr;
@@ -1991,6 +2101,36 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
19912101
inner_dimensions[0] = *count_ptr;
19922102
innerloop(dataptr, inner_dimensions, inner_strides, innerloopdata);
19932103
} while (iternext(iter));
2104+
} else {
2105+
/**
2106+
* For each output operand, check if it has non-zero size,
2107+
* and assign the identity if it does. For example, a dot
2108+
* product of two zero-length arrays will be a scalar,
2109+
* which has size one.
2110+
*/
2111+
for (i = nin; i < nop; ++i) {
2112+
if (PyArray_SIZE(op[i]) != 0) {
2113+
switch (ufunc->identity) {
2114+
case PyUFunc_Zero:
2115+
assign_reduce_identity_zero(op[i]);
2116+
break;
2117+
case PyUFunc_One:
2118+
assign_reduce_identity_one(op[i]);
2119+
break;
2120+
case PyUFunc_None:
2121+
case PyUFunc_ReorderableNone:
2122+
PyErr_Format(PyExc_ValueError,
2123+
"ufunc %s ",
2124+
ufunc_name);
2125+
goto fail;
2126+
default:
2127+
PyErr_Format(PyExc_ValueError,
2128+
"ufunc %s has an invalid identity for reduction",
2129+
ufunc_name);
2130+
goto fail;
2131+
}
2132+
}
2133+
}
19942134
}
19952135

19962136
/* Check whether any errors occurred during the loop */
@@ -2433,13 +2573,13 @@ reduce_type_resolver(PyUFuncObject *ufunc, PyArrayObject *arr,
24332573
}
24342574

24352575
static int
2436-
assign_reduce_identity_zero(PyArrayObject *result, void *data)
2576+
assign_reduce_identity_zero(PyArrayObject *result)
24372577
{
24382578
return PyArray_FillWithScalar(result, PyArrayScalar_False);
24392579
}
24402580

24412581
static int
24 8000 42-
assign_reduce_identity_one(PyArrayObject *result, void *data)
2582+
assign_reduce_identity_one(PyArrayObject *result)
24432583
{
24442584
return PyArray_FillWithScalar(result, PyArrayScalar_True);
24452585
}

‎numpy/core/tests/test_ufunc.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,8 @@ def test_forced_sig(self):
316316
def test_inner1d(self):
317317
a = np.arange(6).reshape((2,3))
318318
assert_array_equal(umt.inner1d(a,a), np.sum(a*a,axis=-1))
319+
a = np.arange(6)
320+
assert_array_equal(umt.inner1d(a,a), np.sum(a*a))
319321

320322
def test_broadcast(self):
321323
msg = "broadcast"
@@ -425,6 +427,13 @@ def test_innerwt(self):
425427
w = np.arange(300,324).reshape((2,3,4))
426428
assert_arra 7AC8 y_equal(umt.innerwt(a,b,w), np.sum(a*b*w,axis=-1))
427429

430+
def test_innerwt_empty(self):
431+
"""Test generalized ufunc with zero-sized operands"""
432+
a = np.array([], dtype='f8')
433+
b = np.array([], dtype='f8')
434+
w = np.array([], dtype='f8')
435+
assert_array_equal(umt.innerwt(a,b,w), np.sum(a*b*w,axis=-1))
436+
428437
def test_matrix_multiply(self):
429438
self.compare_matrix_multiply_results(np.long)
430439
self.compare_matrix_multiply_results(np.double)

0 commit comments

Comments
 (0)
0