TF::PartitionedCallOp call_op) const {
    // If unit-wise quantization config is found, overwrite the default config.
    auto quantization_unit = FindQuantizationUnitFromLoc(call_op.getLoc());
    if (!quantization_unit.has_value()) return std::nullopt;

    for (const auto& unit_config :
         quant_options_.unit_wise_quantization_specs()) {
      for (const auto& unit : unit_config.unit()) {

    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,
                                src_arg.getLoc());

    const std::string new_input_name =
        absl::StrCat(GetInitializerType(src_func_op), "_", src_arg_idx, ":0");

    AddEntryFunctionInput(new_input_name, main_func_op);

  ArrayAttr ids_attr = builder.getStrArrayAttr(ids);
  ArrayAttr calibration_methods_attr =
      builder.getI32ArrayAttr(calibration_methods);
  builder.create<TF::CalibrationStatisticsSaverOp>(
      region.getLoc(), statistics_outputs, output_file_path_attr, ids_attr,
      calibration_methods_attr);
  return success();
}

// Returns true if the op contains a `CalibrationStatisticsSaverOp`.

    if (it.second) {
      auto resource_type = op.getResult(0).getType();
      op.getResult(0).replaceAllUsesWith(
          block.addArgument(resource_type, op.getLoc()));
      AddEntryFunctionInput(
          absl::StrCat("hash_table_", it.first->getValue(), ":0"), target_func);
      // Avoid deleting the op here, clone it to the caller function first.

  TypeRange result_types{ValueRange{results}};
  auto func_type = FunctionType::get(context, arg_types, result_types);

  SmallVector<Location> arg_locs;
  for (Value arg : arguments) {
    arg_locs.push_back(arg.getLoc());
  }

  auto wrap_func = builder.create<func::FuncOp>(location, func_name, func_type);
  wrap_func.setVisibility(SymbolTable::Visibility::Private);

  };
  const SmallVector<Value> operands(dot_general_op->getOperands());
  const SmallVector<Value> results(dot_general_op->getResults());
  Operation* lifted_op =
      LiftAsFunctionCall(builder_, dot_general_op->getLoc(),
                         FunctionCallOpType::TFXlaCallModuleOp,
                         "composite_dot_general_fn", operands, results,
                         attributes)[0]
          .getDefiningOp();

      };

      // Insert custom aggregation op between operand and operator.
      rewriter.setInsertionPointAfterValue(value);
      Operation *aggregator_op = rewriter.create<TF::CustomAggregatorOp>(
          op->getLoc(), output_types, value, attributes);

      Value aggregator_op_result = aggregator_op->getOpResult(0);
      value.replaceAllUsesExcept(aggregator_op_result, aggregator_op);
    }

    return success();
  }

            return !wrapped_ops.count(user.getOwner());
          }))
        result_types.push_back(result.getType());
    }
  }

  IslandOp new_island = OpBuilder(island).create<IslandOp>(
      island.getLoc(), result_types,
      /*control=*/ControlType::get(island.getContext()),
      /*controlInputs=*/island.getOperands());
  new_island.getBody().push_back(new Block);

        TensorShapeProto shape_proto;
        context->ShapeHandleToProto(handle, &shape_proto);
        if (!shape_proto.unknown_rank()) {
          NodeDef const_def;
          const_def.set_op("Const");
          Node* var_node;
          TF_RETURN_IF_ERROR(n->input_node(0, &var_node));
          const_def.set_name(
              graph->NewName(absl::StrCat("var_shape_", var_node->name())));

        }
      }
      continue;
    }

    if (IsSplitSharding(sharding)) {
      llvm::SmallVector<mlir::Value, 4> tiled_inputs;
      auto result = HandleTileShardedInputs(
          cluster_func.getLoc(), sharding, input_value, builder, &tiled_inputs);
      if (mlir::failed(result)) return mlir::failure();

      const int64_t tiled_inputs_size = tiled_inputs.size();
      if (tiled_inputs_size != num_cores_per_replica)