SmallVector<Type> result_types{op.getResult().getType()};
    ReplaceTfOpWithNewOp<TF::XlaVariadicReduceV2Op>(
        rewriter, op, result_types, inputs, init_values,
        op.getDimensionsToReduce(), op.getReducer());
    return ::mlir::success();
  };
};

void XlaReduceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {

    } 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());
    }
  }