argument_indices_to_remove.push_back(index);
    } else {
      if (is_var_handle) {
        // Add resource_name attribute to VarHandleOp read.
        function.setArgAttr(
            new_argument_index, kResourceNameArgAttr,
            builder.getStringAttr(
                var_handle_shared_names[index - var_handles_start_idx]));
      }
      ++new_argument_index;
    }

        }
      } else if (!operand.getParentRegion()->isProperAncestor(
                     &replicate.getBody())) {
        // Not a replication-invariant operand.
        continue;
      }
      func.setArgAttr(entry.index(), kReplicationAttr,
                      builder.getBoolAttr(true));
    }
  });
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>
CreateAnnotateParameterReplicationPass() {

      /*sym_name=*/
      builder.getStringAttr("__tf_saved_model_variables"),  // Val unimportant.
      /*filename=*/
      builder.getStringAttr(checkpoint_path));
  // Marks the input to be inlined.
  main_func.setArgAttr(
      0, "tf_saved_model.bound_input",
      SymbolRefAttr::get(builder.getContext(), asset_op.getName()));
  // Bound arguments are expected to be of type tensor<!tf_type.string>, not
  // tensor<x!tf_type.string>.

        alias_info.input_index, kAliasingAttr);

    // Set only if aliasing attribute does not exist.
    if (!aliasing_attr) {
      device_func.setArgAttr(
          alias_info.input_index, kAliasingAttr,
          builder.getI64IntegerAttr(alias_info.output_index));
      continue;
    }
    // If aliasing attribute already exists, it must match the new value.

      auto flat_symbol_ref_attr =
          FlatSymbolRefAttr::get(context, resource_name);

      // Add the corresponding `tf_saved_model.bound_input` attribute.
      func.setArgAttr(i, kSavedModelArgAttr, flat_symbol_ref_attr);

      resource_names.insert(flat_symbol_ref_attr.getValue());

      // Remove the existing `tf.resource_name` attribute.
      func.removeArgAttr(i, resource_name_id);

      builder.getFunctionType(input_types, main_func_op.getResultTypes()));

  // Add "__tf_file_prefix" to the "tf_saved_model.index_path" attribute for the
  // newly created argument.
  main_func_op.setArgAttr(new_file_prefix_arg.getArgNumber(),
                          /*name=*/kTfSavedModelIndexPathAttr,
                          /*value=*/builder.getStrArrayAttr({kTfFilePrefix}));

    if (device_attr.empty()) {
      // If device_attr does not exist, try to construct it from any recorded
      // assignment.
      if (auto device = result->DeviceForResource(arg)) {
        func_op.setArgAttr(arg.getArgNumber(), kFuncDeviceAttr,
                           builder.getStringAttr(*device));
      }
      continue;
    }
    // Record the attribute.

    // Sets the device attribute for every input and every result of the
    // function.
    for (int i : llvm::seq<int>(0, metadata.input_devices.size())) {
      func_op.setArgAttr(i, kTFDeviceAttr,
                         StringAttr::get(context, metadata.input_devices[i]));
    }
    for (int i : llvm::seq<int>(0, metadata.result_devices.size())) {
      func_op.setResultAttr(

        << ", got: " << main_func.getNumResults();
    return false;
  }

  const int num_args = main_func.getNumArguments();
  for (int i = 0; i < num_args; ++i) {
    main_func.setArgAttr(
        i, kTfSavedModelIndexPathAttr,
        ArrayAttr::get(context, {StringAttr::get(context, input_names[i])}));
  }

  const int num_results = main_func.getNumResults();

    BlockArgument arg = std::get<1>(sharding_and_arg);
    const auto& sharding_or_op = std::get<0>(sharding_and_arg);
    if (auto sharding = GetShardingStringFromVariant(sharding_or_op)) {
      func.setArgAttr(arg.getArgNumber(), kShardingAttr,
                      builder->getStringAttr(*sharding));
    }
  }

  Operation* terminator = function_block.getTerminator();
  for (auto sharding_and_retval :