// `XlaGatherOp (https://www.tensorflow.org/xla/operation_semantics#gather). The
// resulting type's shape has `tensor_type.size() + collapsed_dims.size()`
// dimensions.
RankedTensorType RestoreCollapsedDimensions(
    const RankedTensorType tensor_type,
    const absl::flat_hash_set<int64_t>& collapsed_dims) {
  ArrayRef<int64_t> original_tensor_shape = tensor_type.getShape();
  const int output_tensor_rank =