I don’t think gather can do that since it’s indexing only 1 dim
currently I used reshape to flat the indexing matrices, and reshape the resulted image back to the correct shape… but that seems kind of awkward.