8000 ENH: change list-of-array to tuple-of-array returns (Numba compat) by rgommers · Pull Request #25570 · numpy/numpy · GitHub
[go: up one dir, main page]

Skip to content

ENH: change list-of-array to tuple-of-array returns (Numba compat) #25570

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

Merged
merged 5 commits into from
Jan 21, 2024
Merged
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
7 changes: 7 additions & 0 deletions doc/release/upcoming_changes/25570.change
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
Functions that returned a list of ndarrays have been changed to return a tuple
of ndarrays instead. Returning tuples consistently whenever a sequence of
arrays is returned makes it easier for JIT compilers like Numba, as well as for
static type checkers in some cases, to support these functions. Changed
functions are: ``atleast_1d``, ``atleast_2d``, ``atleast_3d``, ``broadcast_arrays``,
``array_split``, ``split``, ``hsplit``, ``vsplit``, ``dsplit``, ``meshgrid``,
``ogrid``, ``histogramdd``.
24 changes: 12 additions & 12 deletions numpy/_core/shape_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ def atleast_1d(*arys):
Returns
-------
ret : ndarray
An array, or list of arrays, each with ``a.ndim >= 1``.
An array, or tuple of arrays, each with ``a.ndim >= 1``.
Copies are made only if necessary.

See Also
Expand All @@ -57,7 +57,7 @@ def atleast_1d(*arys):
True

>>> np.atleast_1d(1, [3, 4])
[array([1]), array([3, 4])]
(array([1]), array([3, 4]))

"""
res = []
Expand All @@ -71,7 +71,7 @@ def atleast_1d(*arys):
if len(res) == 1:
return res[0]
else:
return res
return tuple(res)


def _atleast_2d_dispatcher(*arys):
Expand All @@ -93,7 +93,7 @@ def atleast_2d(*arys):
Returns
-------
res, res2, ... : ndarray
An array, or list of arrays, each with ``a.ndim >= 2``.
An array, or tuple of arrays, each with ``a.ndim >= 2``.
Copies are avoided where possible, and views with two or more
dimensions are returned.

Expand All @@ -113,7 +113,7 @@ def atleast_2d(*arys):
True

>>> np.atleast_2d(1, [1, 2], [[1, 2]])
[array([[1]]), array([[1, 2]]), array([[1, 2]])]
(array([[1]]), array([[1, 2]]), array([[1, 2]]))

"""
res = []
Expand All @@ -129,7 +129,7 @@ def atleast_2d(*arys):
if len(res) == 1:
return res[0]
else:
return res
return tuple(res)


def _atleast_3d_dispatcher(*arys):
Expand All @@ -151,7 +151,7 @@ def atleast_3d(*arys):
Returns
-------
res1, res2, ... : ndarray
An array, or list of arrays, each with ``a.ndim >= 3``. Copies are
An array, or tuple of arrays, each with ``a.ndim >= 3``. Copies are
avoided where possible, and views with three or more dimensions are
returned. For example, a 1-D array of shape ``(N,)`` becomes a view
of shape ``(1, N, 1)``, and a 2-D array of shape ``(M, N)`` becomes a
Expand Down Expand Up @@ -201,7 +201,7 @@ def atleast_3d(*arys):
if len(res) == 1:
return res[0]
else:
return res
return tuple(res)


def _arrays_for_stack_dispatcher(arrays):
Expand Down Expand Up @@ -282,8 +282,8 @@ def vstack(tup, *, dtype=None, casting="same_kind"):

"""
arrs = atleast_2d(*tup)
if not isinstance(arrs, list):
arrs = [arrs]
if not isinstance(arrs, tuple):
arrs = (arrs,)
return _nx.concatenate(arrs, 0, dtype=dtype, casting=casting)


Expand Down Expand Up @@ -349,8 +349,8 @@ def hstack(tup, *, dtype=None, casting="same_kind"):

"""
arrs = atleast_1d(*tup)
if not isinstance(arrs, list):
arrs = [arrs]
if not isinstance(arrs, tuple):
arrs = (arrs,)
# As a special case, dimension 0 of 1-dimensional arrays is "horizontal"
if arrs and arrs[0].ndim == 1:
return _nx.concatenate(arrs, 0, dtype=dtype, casting=casting)
Expand Down
6 changes: 3 additions & 3 deletions numpy/_core/shape_base.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -20,21 +20,21 @@ def atleast_1d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
@overload
def atleast_1d(arys: ArrayLike, /) -> NDArray[Any]: ...
@overload
def atleast_1d(*arys: ArrayLike) -> list[NDArray[Any]]: ...
def atleast_1d(*arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...

@overload
def atleast_2d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
@overload
def atleast_2d(arys: ArrayLike, /) -> NDArray[Any]: ...
@overload
def atleast_2d(*arys: ArrayLike) -> list[NDArray[Any]]: ...
def atleast_2d(*arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...

@overload
def atleast_3d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
@overload
def atleast_3d(arys: ArrayLike, /) -> NDArray[Any]: ...
@overload
def atleast_3d(*arys: ArrayLike) -> list[NDArray[Any]]: ...
def atleast_3d(*arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...

@overload
def vstack(
Expand Down
6 changes: 3 additions & 3 deletions numpy/lib/_function_base_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -4986,7 +4986,7 @@ def _meshgrid_dispatcher(*xi, copy=None, sparse=None, indexing=None):
@array_function_dispatch(_meshgrid_dispatcher)
def meshgrid(*xi, copy=True, sparse=False, indexing='xy'):
"""
Return a list of coordinate matrices from coordinate vectors.
Return a tuple of coordinate matrices from coordinate vectors.

Make N-D coordinate arrays for vectorized evaluations of
N-D scalar/vector fields over N-D grids, given
Expand Down Expand Up @@ -5027,7 +5027,7 @@ def meshgrid(*xi, copy=True, sparse=False, indexing='xy'):

Returns
-------
X1, X2,..., XN : list of ndarrays
X1, X2,..., XN : tuple of ndarrays
For vectors `x1`, `x2`,..., `xn` with lengths ``Ni=len(xi)``,
returns ``(N1, N2, N3,..., Nn)`` shaped arrays if indexing='ij'
or ``(N2, N1, N3,..., Nn)`` shaped arrays if indexing='xy'
Expand Down Expand Up @@ -5136,7 +5136,7 @@ def meshgrid(*xi, copy=True, sparse=False, indexing='xy'):
output = np.broadcast_arrays(*output, subok=True)

if copy:
output = [x.copy() for x in output]
output = tuple(x.copy() for x in output)

return output

Expand Down
2 changes: 1 addition & 1 deletion numpy/lib/_function_base_impl.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -626,7 +626,7 @@ def meshgrid(
copy: bool = ...,
sparse: bool = ...,
indexing: L["xy", "ij"] = ...,
) -> list[NDArray[Any]]: ...
) -> tuple[NDArray[Any], ...]: ...

@overload
def delete(
Expand Down
4 changes: 2 additions & 2 deletions numpy/lib/_histograms_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -957,8 +957,8 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None):
H : ndarray
The multidimensional histogram of sample x. See density and weights
for the different possible semantics.
edges : list
A list of D arrays describing the bin edges for each dimension.
edges : tuple of ndarrays
A tuple of D arrays describing the bin edges for each dimension.

See Also
--------
Expand Down
2 changes: 1 addition & 1 deletion numpy/lib/_histograms_impl.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -44,4 +44,4 @@ def histogramdd(
range: Sequence[tuple[float, float]] = ...,
density: None | bool = ...,
weights: None | ArrayLike = ...,
) -> tuple[NDArray[Any], list[NDArray[Any]]]: ...
) -> tuple[NDArray[Any], tuple[NDArray[Any], ...]]: ...
28 changes: 20 additions & 8 deletions numpy/lib/_index_tricks_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,8 @@ def __getitem__(self, key):
slobj[k] = slice(None, None)
nn[k] = nn[k][tuple(slobj)]
slobj[k] = _nx.newaxis
return nn
return tuple(nn) # ogrid -> tuple of arrays
return nn # mgrid -> ndarray
except (IndexError, TypeError):
step = key.step
stop = key.stop
Expand Down Expand Up @@ -225,8 +226,9 @@ class MGridClass(nd_grid):

Returns
-------
mesh-grid
`ndarray`\\ s all of the same dimensions
mesh-grid : ndarray
A single array, containing a set of `ndarray`\\ s all of the same
dimensions. stacked along the first axis.

See Also
--------
Expand All @@ -251,6 +253,13 @@ class MGridClass(nd_grid):
>>> np.mgrid[-1:1:5j]
array([-1. , -0.5, 0. , 0.5, 1. ])

>>> np.mgrid[0:4].shape
(4,)
>>> np.mgrid[0:4, 0:5].shape
(2, 4, 5)
>>> np.mgrid[0:4, 0:5, 0:6].shape
(3, 4, 5, 6)

"""

def __init__(self):
Expand All @@ -277,8 +286,10 @@ class OGridClass(nd_grid):

Returns
-------
mesh-grid
`ndarray`\\ s with only one dimension not equal to 1
mesh-grid : ndarray or tuple of ndarrays
If the input is a single slice, returns an array.
If the input is multiple slices, returns a tuple of arrays, with
only one dimension not equal to 1.

See Also
--------
Expand All @@ -292,12 +303,13 @@ class OGridClass(nd_grid):
>>> from numpy import ogrid
>>> ogrid[-1:1:5j]
array([-1. , -0.5, 0. , 0.5, 1. ])
>>> ogrid[0:5,0:5]
[array([[0],
>>> ogrid[0:5, 0:5]
(array([[0],
[1],
[2],
[3],
[4]]), array([[0, 1, 2, 3, 4]])]
[4]]),
array([[0, 1, 2, 3, 4]]))

"""

Expand Down
2 changes: 1 addition & 1 deletion numpy/lib/_index_tricks_impl.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ class nd_grid(Generic[_BoolType]):
def __getitem__(
self: nd_grid[Literal[True]],
key: slice | Sequence[slice],
) -> list[NDArray[Any]]: ...
) -> tuple[NDArray[Any], ...]: ...

class MGridClass(nd_grid[Literal[False]]):
def __init__(self) -> None: ...
Expand Down
Loading
0