summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
authorHuang, Guangtai <guangtai@amazon.com>2019-08-12 12:34:01 +0800
committerHuang, Guangtai <guangtai@amazon.com>2019-08-13 11:43:20 +0800
commit3b3eaa60dc31677d1b19cb49e5c414d41ac4089a (patch)
tree65df45e7f76fa9bdae959cae8343fea8deb8fbd8 /numpy
parentcff02456fed4113924feef1d561e5e927aa2d86b (diff)
downloadnumpy-3b3eaa60dc31677d1b19cb49e5c414d41ac4089a.tar.gz
fix unique
Diffstat (limited to 'numpy')
-rw-r--r--numpy/lib/arraysetops.py4
-rw-r--r--numpy/lib/tests/test_arraysetops.py7
2 files changed, 7 insertions, 4 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py
index f3f4bc17e..c3e833f74 100644
--- a/numpy/lib/arraysetops.py
+++ b/numpy/lib/arraysetops.py
@@ -264,7 +264,7 @@ def unique(ar, return_index=False, return_inverse=False,
# axis was specified and not None
try:
- ar = np.swapaxes(ar, axis, 0)
+ ar = np.moveaxis(ar, axis, 0)
except np.AxisError:
# this removes the "axis1" or "axis2" prefix from the error message
raise np.AxisError(axis, ar.ndim)
@@ -285,7 +285,7 @@ def unique(ar, return_index=False, return_inverse=False,
def reshape_uniq(uniq):
uniq = uniq.view(orig_dtype)
uniq = uniq.reshape(-1, *orig_shape[1:])
- uniq = np.swapaxes(uniq, 0, axis)
+ uniq = np.moveaxis(uniq, 0, axis)
return uniq
output = _unique1d(consolidated, return_index,
diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py
index dd8a38248..fd21a7f76 100644
--- a/numpy/lib/tests/test_arraysetops.py
+++ b/numpy/lib/tests/test_arraysetops.py
@@ -600,8 +600,11 @@ class TestUnique(object):
assert_array_equal(unique(data, axis=1), result.astype(dtype), msg)
msg = 'Unique with 3d array and axis=2 failed'
- data3d = np.dstack([data] * 3)
- result = data3d[..., :1]
+ data3d = np.array([[[1, 1],
+ [1, 0]],
+ [[0, 1],
+ [0, 0]]]).astype(dtype)
+ result = np.take(data3d, [1, 0], axis=2)
assert_array_equal(unique(data3d, axis=2), result, msg)
uniq, idx, inv, cnt = unique(data, axis=0, return_index=True,