} else if (auto xla_variadic_reduce_v2_op =
                   dyn_cast<TF::XlaVariadicReduceV2Op>(op)) {
      return propagate_shape_to(xla_variadic_reduce_v2_op.getReducer());
    } else if (auto xla_variadic_sort_op =
                   dyn_cast<TF::XlaVariadicSortOp>(op)) {
      return propagate_shape_to(xla_variadic_sort_op.getComparator());
    }
  }

  LogicalResult matchAndRewrite(TF::XlaVariadicReduceV2Op op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();

    mlir::SymbolRefAttr func = op.getReducer();
    auto func_op = cast<mlir::func::FuncOp>(SymbolTable::lookupSymbolIn(
        op->getParentOfType<mlir::ModuleOp>(), func));
    auto func_ty = func_op.getFunctionType();