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

`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