Issue
There is np.stack in NumPy, but is there an opposite np.unstack
same as tf.unstack?
Solution
Coming across this late, here is a much simpler answer:
def unstack(a, axis=0):
return np.moveaxis(a, axis, 0)
# return list(np.moveaxis(a, axis, 0))
As a bonus, the result is still a numpy array. The unwrapping happens if you just python-unwrap it:
A, B, = unstack([[1, 2], [3, 4]], axis=1)
assert list(A) == [1, 3]
assert list(B) == [2, 4]
Unsurprisingly, it is also the fastest:
# np.squeeze
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = [np.squeeze(e, 1) for e in np.split(a, a.shape[1], axis = 1)]"
100 loops, best of 5: 2.64 msec per loop
# np.take
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = [np.take(a, i, axis = 1) for i in range(a.shape[1])]"
50 loops, best of 5: 5.08 msec per loop
# np.moveaxis
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = np.moveaxis(a, 1, 0)"
100000 loops, best of 5: 3.89 usec per loop
# list(np.moveaxis)
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = list(np.moveaxis(a, 1, 0))"
1000 loops, best of 5: 205 usec per loop
Answered By - Ivorius
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.