From b6313fdaa4d129e4b58684b0bede9696edde92d5 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 15 Jan 2021 20:18:47 +0100 Subject: import indices, broadcast_to --- numpy/lib/twodim_base.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) (limited to 'numpy/lib') diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py index 05e92b590..960797b68 100644 --- a/numpy/lib/twodim_base.py +++ b/numpy/lib/twodim_base.py @@ -6,11 +6,12 @@ import functools from numpy.core.numeric import ( asanyarray, arange, zeros, greater_equal, multiply, ones, asarray, where, int8, int16, int32, int64, empty, promote_types, diagonal, - nonzero + nonzero, indices ) from numpy.core.overrides import set_array_function_like_doc, set_module from numpy.core import overrides from numpy.core import iinfo +from numpy.lib.stride_tricks import broadcast_to __all__ = [ @@ -894,10 +895,10 @@ def tril_indices(n, k=0, m=None): [-10, -10, -10, -10]]) """ - tri = np.tri(n, m, k=k, dtype=bool) + tri_ = tri(n, m, k=k, dtype=bool) - return tuple(np.broadcast_to(inds, tri.shape)[tri] - for inds in np.indices(tri.shape, sparse=True)) + return tuple(broadcast_to(inds, tri_.shape)[tri_] + for inds in indices(tri_.shape, sparse=True)) def _trilu_indices_form_dispatcher(arr, k=None): @@ -1013,10 +1014,10 @@ def triu_indices(n, k=0, m=None): [ 12, 13, 14, -1]]) """ - tri = ~np.tri(n, m, k=k - 1, dtype=bool) + tri_ = ~tri(n, m, k=k - 1, dtype=bool) - return tuple(np.broadcast_to(inds, tri.shape)[tri] - for inds in np.indices(tri.shape, sparse=True)) + return tuple(broadcast_to(inds, tri_.shape)[tri_] + for inds in indices(tri_.shape, sparse=True)) @array_function_dispatch(_trilu_indices_form_dispatcher) -- cgit v1.2.1