if (!symbol) return failure();
    auto f = llvm::dyn_cast<mlir::func::FuncOp>(symbol);

    if (call.getOperation()->getNumOperands() !=
        f.getBody().front().getNumArguments()) {
      return failure();  // RemoteCall et al
    }

    rewriter.startOpModification(f);
    bool changed = false;
    for (auto [call_arg, body_arg] :

  // all those post are loop independent. Initialize extern_values with while_op
  // not loop carried operands.
  auto num_loop_carried = while_op.getCond().getNumArguments();
  auto not_carried_operands =
      while_op.getOperands().drop_front(num_loop_carried);
  extern_values.insert(not_carried_operands.begin(),
                       not_carried_operands.end());

    entry->addArguments(inputs, SmallVector<Location>(inputs.size(), loc));
    args = entry->getArguments();
  }
  llvm::SmallVector<Value, 4> casted_args;
  casted_args.reserve(func.getNumArguments());
  for (const auto& ArgAndType : zip(args, func.getFunctionType().getInputs())) {
    Value arg = std::get<0>(ArgAndType);
    Type expected_type = std::get<1>(ArgAndType);
    if (arg.getType() != expected_type) {

    return success(IsMaxFunction(op.getBody().front()));
  }

  // Whether `block` semantically corresponds to a `stablehlo.maximum` op.
  bool IsMaxFunction(Block& block) const {
    if (block.getNumArguments() != 2) return false;

    auto return_op = cast<mlir::stablehlo::ReturnOp>(block.getTerminator());
    if (return_op.getNumOperands() != 1) return false;

    auto max_op = dyn_cast_or_null<MaxOp>(

// "__tf_file_prefix". If such an argument doesn't exist, returns a null value.
BlockArgument GetFilePrefixArg(func::FuncOp main_func_op) {
  for (int i = 0; i < main_func_op.getNumArguments(); i++) {
    auto index_path_attr =
        main_func_op.getArgAttrOfType<ArrayAttr>(i, kTfSavedModelIndexPathAttr);
    if (index_path_attr && !index_path_attr.empty() &&

        /*tpu_core=*/rewriter.getI64IntegerAttr(0));
  }
};

void UpdateArgAttributes(mlir::func::FuncOp func) {
  OpBuilder builder(func.getBody());
  for (int i = 0; i < func.getNumArguments(); ++i) {
    constexpr char kShardingAttr[] = "mhlo.sharding";
    if (auto sharding =
            func.getArgAttrOfType<mlir::StringAttr>(i, kShardingAttr)) {
      if (!sharding.getValue().empty()) {

Value UpdateControlFlowBlockArgWithToken(OpBuilder& builder, Block& block,
                                         ArrayRef<Type> types) {
  builder.setInsertionPointToStart(&block);

  auto old_args_size = block.getNumArguments();

  block.addArguments(
      types, SmallVector<Location>(types.size(), block.getParent()->getLoc()));

  ValueRange old_args = block.getArguments().take_front(old_args_size);

    OptionalOpShardingVector& sharding_for_args) {
  // Look up function definition from module.
  Block& function_block = func.front();

  sharding_for_args.reserve(function_block.getNumArguments());

  // Iterate through operands of `cluster_func`.
  // The computation operand can either be:
  //   1) a TPUPartitionedInput Op if the input has a non-resource type;
  //   2) a ReadVariableOp else.
  //

    // No need to create a mapping when there is no usage - it will not affect
    // the cloning.
    if (src_arg.use_empty()) continue;

    const unsigned main_arg_idx = main_func_op.getNumArguments();

    const DictionaryAttr main_arg_attr =
        src_func_op.getArgAttrDict(src_arg_idx);

    main_func_op.insertArgument(main_arg_idx, src_arg.getType(), main_arg_attr,

  constexpr char kAliasingAttr[] = "tf.aliasing_output";
  llvm::SmallDenseMap<unsigned, unsigned> output_to_input_alias;
  unsigned num_arguments = main_fn.getNumArguments();
  for (unsigned arg_index = 0; arg_index < num_arguments; ++arg_index) {
    if (auto aliasing_output = main_fn.getArgAttrOfType<mlir::IntegerAttr>(
            arg_index, kAliasingAttr))