Pick vs Take

1

Can someone explain the difference between NDArray.pick and NDArray.take?

2

NDArray.pick picks from the ndarray slices along the dimension specified by axis, using the index specified by indices, while NDArray.take takes the entire ndarray slice and axis is always 0.

Here’s an example that illustrates the difference:

>>> from mxnet import nd

>>> a = nd.array([[1, 2], [3, 4], [5, 6]])

>>> a

[[1. 2.]

 [3. 4.]

 [5. 6.]]

>>> y = nd.array([0, 1])

>>> nd.take(a, y)

[[1. 2.]

 [3. 4.]]

<NDArray 2x2 @cpu(0)>

>>> nd.pick(a, y, 0)

[1. 4.]

<NDArray 2 @cpu(0)>

>>>

Take directly indexes the 0 axis while using indices, while Pick indexes the axis provided using both the index of the index and the index in indices. The documentation page linked earlier has more details.

1 Like