Function take_along_axis 

pub fn take_along_axis(a: &Array, indices: &Array, axis: impl Axis) -> Array

The elements at indices along axis; see Array::take_along_axis.