Skip to content

ENH: use __skip_array_function__ internally in NumPy #13585

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 21 additions & 14 deletions numpy/core/numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,10 +156,11 @@ def zeros_like(a, dtype=None, order='K', subok=True, shape=None):
array([0., 0., 0.])

"""
res = empty_like(a, dtype=dtype, order=order, subok=subok, shape=shape)
res = empty_like.__skip_array_function__(
a, dtype=dtype, order=order, subok=subok, shape=shape)
# needed instead of a 0 to get same result as zeros for for string dtypes
z = zeros(1, dtype=res.dtype)
multiarray.copyto(res, z, casting='unsafe')
multiarray.copyto.__skip_array_function__(res, z, casting='unsafe')
return res


Expand Down Expand Up @@ -212,7 +213,7 @@ def ones(shape, dtype=None, order='C'):

"""
a = empty(shape, dtype, order)
multiarray.copyto(a, 1, casting='unsafe')
multiarray.copyto.__skip_array_function__(a, 1, casting='unsafe')
return a


Expand Down Expand Up @@ -282,8 +283,9 @@ def ones_like(a, dtype=None, order='K', subok=True, shape=None):
array([1., 1., 1.])

"""
res = empty_like(a, dtype=dtype, order=order, subok=subok, shape=shape)
multiarray.copyto(res, 1, casting='unsafe')
res = empty_like.__skip_array_function__(
a, dtype=dtype, order=order, subok=subok, shape=shape)
multiarray.copyto.__skip_array_function__(res, 1, casting='unsafe')
return res


Expand Down Expand Up @@ -330,7 +332,7 @@ def full(shape, fill_value, dtype=None, order='C'):
if dtype is None:
dtype = array(fill_value).dtype
a = empty(shape, dtype, order)
multiarray.copyto(a, fill_value, casting='unsafe')
multiarray.copyto.__skip_array_function__(a, fill_value, casting='unsafe')
return a


Expand Down Expand Up @@ -397,8 +399,10 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):
array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1])

"""
res = empty_like(a, dtype=dtype, order=order, subok=subok, shape=shape)
multiarray.copyto(res, fill_value, casting='unsafe')
res = empty_like.__skip_array_function__(
a, dtype=dtype, order=order, subok=subok, shape=shape)
multiarray.copyto.__skip_array_function__(
res, fill_value, casting='unsafe')
return res


Expand Down Expand Up @@ -575,7 +579,8 @@ def argwhere(a):
[1, 2]])

"""
return transpose(nonzero(a))
return transpose.__skip_array_function__(
nonzero.__skip_array_function__(a))


def _flatnonzero_dispatcher(a):
Expand Down Expand Up @@ -620,7 +625,8 @@ def flatnonzero(a):
array([-2, -1, 1, 2])

"""
return np.nonzero(np.ravel(a))[0]
return np.nonzero.__skip_array_function__(
np.ravel.__skip_array_function__(a))[0]


_mode_from_name_dict = {'v': 0,
Expand Down Expand Up @@ -2127,7 +2133,8 @@ def allclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
True

"""
res = all(isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan))
res = all(isclose.__skip_array_function__(
a, b, rtol=rtol, atol=atol, equal_nan=equal_nan))
return bool(res)


Expand Down Expand Up @@ -2231,12 +2238,12 @@ def within_tol(x, y, atol, rtol):
return within_tol(x, y, atol, rtol)
else:
finite = xfin & yfin
cond = zeros_like(finite, subok=True)
cond = zeros_like.__skip_array_function__(finite, subok=True)
# Because we're using boolean indexing, x & y must be the same shape.
# Ideally, we'd just do x, y = broadcast_arrays(x, y). It's in
# lib.stride_tricks, though, so we can't import it here.
x = x * ones_like(cond)
y = y * ones_like(cond)
x = x * ones_like.__skip_array_function__(cond)
y = y * ones_like.__skip_array_function__(cond)
# Avoid subtraction with infinite/nan values...
cond[finite] = within_tol(x[finite], y[finite], atol, rtol)
# Check for equality of infinite values...
Expand Down
21 changes: 12 additions & 9 deletions numpy/core/shape_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -273,7 +273,8 @@ def vstack(tup):
[4]])

"""
return _nx.concatenate([atleast_2d(_m) for _m in tup], 0)
return _nx.concatenate.__skip_array_function__(
[atleast_2d.__skip_array_function__(_m) for _m in tup], 0)


@array_function_dispatch(_vhstack_dispatcher)
Expand Down Expand Up @@ -324,12 +325,12 @@ def hstack(tup):
[3, 4]])

"""
arrs = [atleast_1d(_m) for _m in tup]
arrs = [atleast_1d.__skip_array_function__(_m) for _m in tup]
# As a special case, dimension 0 of 1-dimensional arrays is "horizontal"
if arrs and arrs[0].ndim == 1:
return _nx.concatenate(arrs, 0)
return _nx.concatenate.__skip_array_function__(arrs, 0)
else:
return _nx.concatenate(arrs, 1)
return _nx.concatenate.__skip_array_function__(arrs, 1)


def _stack_dispatcher(arrays, axis=None, out=None):
Expand Down Expand Up @@ -413,7 +414,8 @@ def stack(arrays, axis=0, out=None):

sl = (slice(None),) * axis + (_nx.newaxis,)
expanded_arrays = [arr[sl] for arr in arrays]
return _nx.concatenate(expanded_arrays, axis=axis, out=out)
return _nx.concatenate.__skip_array_function__(
expanded_arrays, axis=axis, out=out)


def _block_format_index(index):
Expand Down Expand Up @@ -496,8 +498,8 @@ def _block_check_depths_match(arrays, parent_index=[]):
return parent_index + [None], 0, 0
else:
# We've 'bottomed out' - arrays is either a scalar or an array
size = _nx.size(arrays)
return parent_index, _nx.ndim(arrays), size
size = _nx.size.__skip_array_function__(arrays)
return parent_index, _nx.ndim.__skip_array_function__(arrays), size


def _atleast_nd(a, ndim):
Expand Down Expand Up @@ -640,7 +642,8 @@ def _block(arrays, max_depth, result_ndim, depth=0):
if depth < max_depth:
arrs = [_block(arr, max_depth, result_ndim, depth+1)
for arr in arrays]
return _nx.concatenate(arrs, axis=-(max_depth-depth))
return _nx.concatenate.__skip_array_function__(
arrs, axis=-(max_depth-depth))
else:
# We've 'bottomed out' - arrays is either a scalar or an array
# type(arrays) is not list
Expand Down Expand Up @@ -858,7 +861,7 @@ def _block_slicing(arrays, list_ndim, result_ndim):

# Test preferring F only in the case that all input arrays are F
F_order = all(arr.flags['F_CONTIGUOUS'] for arr in arrays)
C_order = all(arr.flags['C_CONTIGUOUS'] for arr in arrays)
C_order = all(arr.flags['C_CONTIGUOUS'] for arr in arrays)
order = 'F' if F_order and not C_order else 'C'
result = _nx.empty(shape=shape, dtype=dtype, order=order)
# Note: In a c implementation, the function
Expand Down
2 changes: 1 addition & 1 deletion numpy/lib/arraypad.py
Original file line number Diff line number Diff line change
Expand Up @@ -804,7 +804,7 @@ def pad(array, pad_width, mode='constant', **kwargs):
# function operates inplace on the padded array.

# view with the iteration axis at the end
view = np.moveaxis(padded, axis, -1)
view = np.moveaxis.__skip_array_function__(padded, axis, -1)

# compute indices for the iteration axes, and append a trailing
# ellipsis to prevent 0d arrays decaying to scalars (gh-8642)
Expand Down
47 changes: 29 additions & 18 deletions numpy/lib/function_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,26 +134,29 @@ def rot90(m, k=1, axes=(0,1)):
raise ValueError("Axes must be different.")

if (axes[0] >= m.ndim or axes[0] < -m.ndim
or axes[1] >= m.ndim or axes[1] < -m.ndim):
or axes[1] >= m.ndim or axes[1] < -m.ndim):
raise ValueError("Axes={} out of range for array of ndim={}."
.format(axes, m.ndim))
.format(axes, m.ndim))

k %= 4

if k == 0:
return m[:]
if k == 2:
return flip(flip(m, axes[0]), axes[1])
return flip.__skip_array_function__(
flip.__skip_array_function__(m, axes[0]), axes[1])

axes_list = arange(0, m.ndim)
(axes_list[axes[0]], axes_list[axes[1]]) = (axes_list[axes[1]],
axes_list[axes[0]])

if k == 1:
return transpose(flip(m,axes[1]), axes_list)
return transpose.__skip_array_function__(
flip.__skip_array_function__(m, axes[1]), axes_list)
else:
# k == 3
return flip(transpose(m, axes_list), axes[1])
return flip.__skip_array_function__(
transpose.__skip_array_function__(m, axes_list), axes[1])


def _flip_dispatcher(m, axis=None):
Expand Down Expand Up @@ -393,9 +396,11 @@ def average(a, axis=None, weights=None, returned=False):
wgt = np.asanyarray(weights)

if issubclass(a.dtype.type, (np.integer, np.bool_)):
result_dtype = np.result_type(a.dtype, wgt.dtype, 'f8')
result_dtype = np.result_type.__skip_array_function__(
a.dtype, wgt.dtype, 'f8')
else:
result_dtype = np.result_type(a.dtype, wgt.dtype)
result_dtype = np.result_type.__skip_array_function__(
a.dtype, wgt.dtype)

# Sanity checks
if a.shape != wgt.shape:
Expand Down Expand Up @@ -606,7 +611,8 @@ def piecewise(x, condlist, funclist, *args, **kw):

if n == n2 - 1: # compute the "otherwise" condition.
condelse = ~np.any(condlist, axis=0, keepdims=True)
condlist = np.concatenate([condlist, condelse], axis=0)
condlist = np.concatenate.__skip_array_function__(
[condlist, condelse], axis=0)
n += 1
elif n != n2:
raise ValueError(
Expand Down Expand Up @@ -690,13 +696,13 @@ def select(condlist, choicelist, default=0):

# need to get the result type before broadcasting for correct scalar
# behaviour
dtype = np.result_type(*choicelist)
dtype = np.result_type.__skip_array_function__(*choicelist)

# Convert conditions to arrays and broadcast conditions and choices
# as the shape is needed for the result. Doing it separately optimizes
# for example when all choices are scalars.
condlist = np.broadcast_arrays(*condlist)
choicelist = np.broadcast_arrays(*choicelist)
condlist = np.broadcast_arrays.__skip_array_function__(*condlist)
choicelist = np.broadcast_arrays.__skip_array_function__(*choicelist)

# If cond array is not an ndarray in boolean format or scalar bool, abort.
deprecated_ints = False
Expand Down Expand Up @@ -733,7 +739,7 @@ def select(condlist, choicelist, default=0):
choicelist = choicelist[-2::-1]
condlist = condlist[::-1]
for choice, cond in zip(choicelist, condlist):
np.copyto(result, choice, where=cond)
np.copyto.__skip_array_function__(result, choice, where=cond)

return result

Expand Down Expand Up @@ -1512,14 +1518,16 @@ def unwrap(p, discont=pi, axis=-1):
"""
p = asarray(p)
nd = p.ndim
dd = diff(p, axis=axis)
dd = diff.__skip_array_function__(p, axis=axis)
slice1 = [slice(None, None)]*nd # full slices
slice1[axis] = slice(1, None)
slice1 = tuple(slice1)
ddmod = mod(dd + pi, 2*pi) - pi
_nx.copyto(ddmod, pi, where=(ddmod == -pi) & (dd > 0))
_nx.copyto.__skip_array_function__(
ddmod, pi, where=(ddmod == -pi) & (dd > 0))
ph_correct = ddmod - dd
_nx.copyto(ph_correct, 0, where=abs(dd) < discont)
_nx.copyto.__skip_array_function__(
ph_correct, 0, where=abs(dd) < discont)
up = array(p, copy=True, dtype='d')
up[slice1] = p[slice1] + ph_correct.cumsum(axis)
return up
Expand Down Expand Up @@ -1674,7 +1682,10 @@ def extract(condition, arr):
array([0, 3, 6, 9])

"""
return _nx.take(ravel(arr), nonzero(ravel(condition))[0])
return _nx.take.__skip_array_function__(
ravel.__skip_array_function__(arr),
nonzero.__skip_array_function__(
ravel.__skip_array_function__(condition))[0])


def _place_dispatcher(arr, mask, vals):
Expand Down Expand Up @@ -4697,9 +4708,9 @@ def append(arr, values, axis=None):
if axis is None:
if arr.ndim != 1:
arr = arr.ravel()
values = ravel(values)
values = ravel.__skip_array_function__(values)
axis = arr.ndim-1
return concatenate((arr, values), axis=axis)
return concatenate.__skip_array_function__((arr, values), axis=axis)


def _digitize_dispatcher(x, bins, right=None):
Expand Down
21 changes: 12 additions & 9 deletions numpy/lib/histograms.py
Original file line number Diff line number Diff line change
Expand Up @@ -448,7 +448,7 @@ def _search_sorted_inclusive(a, v):

In the context of a histogram, this makes the last bin edge inclusive
"""
return np.concatenate((
return np.concatenate.__skip_array_function__((
a.searchsorted(v[:-1], 'left'),
a.searchsorted(v[-1:], 'right')
))
Expand Down Expand Up @@ -865,17 +865,18 @@ def histogram(a, bins=10, range=None, normed=None, weights=None,
cum_n = np.zeros(bin_edges.shape, ntype)
if weights is None:
for i in _range(0, len(a), BLOCK):
sa = np.sort(a[i:i+BLOCK])
sa = np.sort.__skip_array_function__(a[i:i+BLOCK])
cum_n += _search_sorted_inclusive(sa, bin_edges)
else:
zero = np.zeros(1, dtype=ntype)
for i in _range(0, len(a), BLOCK):
tmp_a = a[i:i+BLOCK]
tmp_w = weights[i:i+BLOCK]
sorting_index = np.argsort(tmp_a)
sorting_index = np.argsort.__skip_array_function__(tmp_a)
sa = tmp_a[sorting_index]
sw = tmp_w[sorting_index]
cw = np.concatenate((zero, sw.cumsum()))
cw = np.concatenate.__skip_array_function__(
(zero, sw.cumsum()))
bin_index = _search_sorted_inclusive(sa, bin_edges)
cum_n += cw[bin_index]

Expand Down Expand Up @@ -1000,7 +1001,7 @@ def histogramdd(sample, bins=10, range=None, normed=None, weights=None,
N, D = sample.shape
except (AttributeError, ValueError):
# Sample is a sequence of 1D arrays.
sample = np.atleast_2d(sample).T
sample = np.atleast_2d.__skip_array_function__(sample).T
N, D = sample.shape

nbin = np.empty(D, int)
Expand All @@ -1027,12 +1028,13 @@ def histogramdd(sample, bins=10, range=None, normed=None, weights=None,

# Create edge arrays
for i in _range(D):
if np.ndim(bins[i]) == 0:
if np.ndim.__skip_array_function__(bins[i]) == 0:
if bins[i] < 1:
raise ValueError(
'`bins[{}]` must be positive, when an integer'.format(i))
smin, smax = _get_outer_edges(sample[:,i], range[i])
edges[i] = np.linspace(smin, smax, bins[i] + 1)
edges[i] = np.linspace.__skip_array_function__(
smin, smax, bins[i] + 1)
elif np.ndim(bins[i]) == 1:
edges[i] = np.asarray(bins[i])
if np.any(edges[i][:-1] > edges[i][1:]):
Expand All @@ -1044,12 +1046,13 @@ def histogramdd(sample, bins=10, range=None, normed=None, weights=None,
'`bins[{}]` must be a scalar or 1d array'.format(i))

nbin[i] = len(edges[i]) + 1 # includes an outlier on each end
dedges[i] = np.diff(edges[i])
dedges[i] = np.diff.__skip_array_function__(edges[i])

# Compute the bin number each sample falls into.
Ncount = tuple(
# avoid np.digitize to work around gh-11022
np.searchsorted(edges[i], sample[:, i], side='right')
np.searchsorted.__skip_array_function__(
edges[i], sample[:, i], side='right')
for i in _range(D)
)

Expand Down
Loading