// The second conditional evaluation.
  builder.setInsertionPointAfter(start_step_0);
  Operation* cond_caller_1 =
      MakeFuncCaller(builder, orig_while_op->getLoc(), orig_cond_func,
                     loop_operands_1, /*flag_for_inlining=*/false);
  Value C_1 = cond_caller_1->getResults().front();

  builder.setInsertionPointAfter(cond_caller_1);
  Operation* non_tpu_caller_1 =

tf_device::LaunchOp CreateNewHostLaunchOpWithNewResult(
    tf_device::LaunchOp* old_launch_op,
    llvm::SmallVector<Value, 4>& new_launch_op_results) {
  OpBuilder builder(*old_launch_op);

  builder.setInsertionPointAfter(*old_launch_op);

  llvm::SmallVector<Type, 4> new_launch_op_results_types;
  for (Value result : new_launch_op_results)
    new_launch_op_results_types.push_back(result.getType());

      if (convert_op && convert_op.getType() == new_result_type) {
        quantizable_op->setOperand(quantize_operand_num, convert_op);
        return;
      }
    }
    rewriter.setInsertionPointAfter(op);
    ConvertOp new_convert_op = rewriter.create<ConvertOp>(
        op->getLoc(), new_result_type, op.getResult());
    quantizable_op->setOperand(quantize_operand_num,

  builder.setInsertionPointAfter(*non_tpu_ops.begin());
  Operation* non_tpu_caller = nullptr;
  result =
      ExtractOpsAsFunc(builder, module, non_tpu_ops, replication_attr, nullptr,
                       nullptr, loop_body_func, "non_tpu", &non_tpu_caller);
  if (failed(result)) return signalPassFailure();

  builder.setInsertionPointAfter(non_tpu_caller);
  Operation* forward_caller = nullptr;

         llvm::maximum(min_max_values[1], min_max_values[3])});
    mlir::ElementsAttr axis_stats;
    mlir::IntegerAttr axis;
    for (auto& stats_op : stats_ops) {
      rewriter.setInsertionPointAfter(stats_op);
      rewriter.replaceOpWithNewOp<quantfork::StatisticsOp>(
          stats_op, stats_op.getArg(), layer_stats, axis_stats, axis);
    }
    return success();
  }
};

  Operation* input_op = input_val.getDefiningOp();

  Operation* insertion_point = input_op->getParentOfType<func::FuncOp>();
  if (!insertion_point) insertion_point = input_op->getParentOfType<ModuleOp>();
  rewriter.setInsertionPointAfter(insertion_point);

  UnrankedTensorType create_unknown_input_shape =
      CreateUnknownShapeFromElementType(input_val.getType());
  UnrankedTensorType create_unknown_output_shape =

  }

  SmallVector<Attribute> result_values;
  if (failed(TF::EvaluateOperation(op, inputs, result_values))) {
    return failure();
  }

  results.clear();
  builder.setInsertionPointAfter(op);
  for (const auto& result_value : result_values) {
    results.push_back(builder.create<TF::ConstOp>(op->getLoc(), result_value));
  }
  return success();
}

bool IsOperationFoldable(Operation* op) {

      if (cast_op && cast_op.getType() == new_result_type) {
        quantize_op->setOperand(quantize_operand_num, cast_op);
        return;
      }
    }
    rewriter.setInsertionPointAfter(op);
    auto new_cast_op =
        rewriter.create<CastOp>(op->getLoc(), new_result_type, op.getResult());
    quantize_op->setOperand(quantize_operand_num, new_cast_op.getResult());
  }

    auto f_name = f_attr.getValue();
    func::FuncOp float_func =
        dyn_cast<func::FuncOp>(symbol_table.lookup(f_name));
    if (!float_func) {
      return failure();
    }
    rewriter.setInsertionPointAfter(call_op);

    // The TPUPartitionedCall has a TPUOrdinalSelectorOp for its last argument
    // which should be removed. So the replaced PartitionedCall op should keep

  if (!func_op->hasAttr(TF::kFromXlaCallModuleAttrName)) {
    return;
  }

  MLIRContext* context = call_op.getContext();
  OpBuilder builder(context);
  builder.setInsertionPointAfter(call_op);

  IRMapping arg_mapper;
  bool call_op_has_platform_index_arg = call_op.getPlatforms().size() > 1;
  // Add an argument for platform_index. This allows for multiple platforms.