void Select(Tensor out,
            const Tensor in,
            const Tensor indices,
            int axis) {
  matchOrAbort<IndexType>(indices->type());
  
  functional::Shape outShape = out->shape();
  functional::Shape inShape  = in->shape();
  functional::Shape idxShape = indices->shape();
  
  int axisCPU = (int)(axis + functional::Shape::size() - out->shape().size());
  
  // Reduce the problem to beforeAxis x idxShape[axisCPU] x afterAxis.
  int beforeAxis = 1;
  for (int i = 0; i < axisCPU; ++i) {
    beforeAxis *= outShape[i];
  }
  int afterAxis = 1;
  for (int i = axisCPU + 1; i < functional::Shape::size(); ++i) {
    afterAxis *= outShape[i];
  }
  // Stride to use for the beforeAxis dimension in the input and output tensors.
  int inBeforeStride = axisCPU ? inShape.stride(axisCPU - 1) : inShape.elements();
  int outBeforeStride = axisCPU ? outShape.stride(axisCPU - 1) : outShape.elements();

  for (int beforeIdx = 0; beforeIdx < beforeAxis; ++beforeIdx) {
    for (int axisIdx = 0; axisIdx < idxShape[axisCPU]; ++axisIdx) {
      // This is the value to read along.
      int index = indices->data<IndexType>()[axisIdx];
      auto inBase = in->data() + beforeIdx * inBeforeStride + index * afterAxis;
      auto outBase = out->data() + beforeIdx * outBeforeStride + axisIdx * afterAxis;
      std::copy(inBase, inBase + afterAxis, outBase);
    }
  }
}