auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
  if (!main_fn) {
    return absl::InternalError("Could not find main function in MLIR Module.");
  }

  mlir::FunctionType func_type = main_fn.getFunctionType();
  for (auto input_type : func_type.getInputs()) {
    tensorflow::TensorShape tensor_shape;
    xla::Shape xla_shape = xla::TypeToShape(input_type);
    TF_RETURN_IF_ERROR(tensorflow::TensorShape::BuildTensorShape(

  DCOMMENT_OP(op.getOperation(), "Infer shape for if ");
  bool changed = false;
  auto then_results =
      op.ResolveThenFunction(&symbol_table_).getFunctionType().getResults();
  auto else_results =
      op.ResolveElseFunction(&symbol_table_).getFunctionType().getResults();
  for (auto it : llvm::zip(op.getResults(), then_results, else_results)) {
    // If then and else types do not match, skip refinement for that result.

    for (auto ret_op : ret->getOperands()) {
      return_types.push_back(ret_op.getType());
    }
    auto newType =
        FunctionType::get(rewriter.getContext(), argument_types, return_types);
    if (f.getFunctionType() == newType) {
      return failure();
    }
    rewriter.modifyOpInPlace(f, [&] { f.setType(newType); });

    // Adjust the type of the return values callers of the function to

                            ConversionPatternRewriter *rewriter) const {
    auto guard = OpBuilder::InsertionGuard(*rewriter);
    auto inputs = branch_func.getFunctionType().getInputs();
    Block *block = rewriter->createBlock(
        &branch_func.getBody(), branch_func.begin(), inputs,
        SmallVector<Location>(inputs.size(), branch_func.getLoc()));

  auto module = op->getParentOfType<ModuleOp>();
  auto function =
      dyn_cast_or_null<func::FuncOp>(SymbolTable::lookupSymbolIn(module, func));
  FunctionType function_ty = function.getFunctionType();

  for (auto arg_in : llvm::zip(args, function_ty.getInputs())) {
    if (std::get<0>(arg_in).getType() != std::get<1>(arg_in)) {
      // Argument type and input type mismatch.
      return false;
    }