Problem

I have two numpy arrays, A and indices.

A has dimensions m x n x 10000. indices has dimensions m x n x 5 (output from argpartition(A, 5)[:,:,:5]). I would like to get a m x n x 5 array containing the elements of A corresponding to indices.

Attempts

indices = np.array([[[5,4,3,2,1],[1,1,1,1,1],[1,1,1,1,1]],
    [500,400,300,200,100],[100,100,100,100,100],[100,100,100,100,100]])
A = np.reshape(range(2 * 3 * 10000), (2,3,10000))

A[...,indices] # gives an array of size (2,3,2,3,5). I want a subset of these values
np.take(A, indices) # shape is right, but it flattens the array first
np.choose(indices, A) # fails because of shape mismatch. 

Motivation

I'm trying to get the 5 largest values of A[i,j] for each i<m, j<n in sorted order using np.argpartition because the arrays can get fairly large.