return RestoreCollapsedDimensions(ranked_output_type, collapsed_dims);
  }

  return xla_gather_op_output_type;
}

// TODO (b/275225582): Supports Xla Gather op in general case.
bool IsXlaGatherWithoutBatch(Value operand, Value start_indices) {
  auto operand_type = mlir::dyn_cast_or_null<ShapedType>(operand.getType());
  auto start_indices_type =
      mlir::dyn_cast_or_null<ShapedType>(start_indices.getType());