pick with multi-dimensional indices - Githubissues

apache / mxnet

Lightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more

https://mxnet.apache.org

Apache License 2.0

20.75k stars 6.8k forks source link

pick with multi-dimensional indices #17786

Open yifeim opened 4 years ago

yifeim commented 4 years ago

Description

Support pick(data, [[0,1,2],[1,2,3]]) = [[data[0,0], data[0,1], data[0,2], [data[1,1], data[1,2], data[1,3]].
As a side note, add a "fill-value" mode to out-of-bound indices. This seems to be a logical replacement to the original "raise"-mode proposal, which is "not supported yet".
The new functions should also support axis parameter in backward compatible ways.

References

Multi-dimensional take is already implemented as take(data, [[0,3],[-1,-2]]) = [[data[0], data[3]], [data[-1], data[-2]].
Since pick is modeled as batch_take, the extension seems natural.

The expected result for a fill-value mode could be verified as

F.pick(data, indices, mode='fill', fill_value=fv) = F.where(
(indices>=0) * (indices<data.shape_array().slice_axis(axis=0, begin=-1, end=None)),
F.pick(data, indices),
F.ones_like(indices.astype(data.dtype))*fv)

yifeim commented 4 years ago

@haojin2