// argument in the list.
    auto arg = bb.getArgument(0);

    auto remove_quantize_op = [&](QuantizeOp quantize_op) {
      auto quantize_output = quantize_op.getOutput();
      auto quantize_type = quantize_output.getType();
      input_types.push_back(quantize_type);
      auto new_arg = bb.addArgument(quantize_type, loc);
      quantize_output.replaceAllUsesWith(new_arg);

          auto dq_arg_type = dq_op.getArg().getType().cast<TensorType>();
          auto qtype = dq_arg_type.getElementType().cast<QuantizedType>();
          auto scast_op = rewriter.create<quantfork::StorageCastOp>(
              dq_op->getLoc(), dq_arg_type.clone(qtype.getStorageType()),
              dq_op.getArg());
          inputs.push_back(scast_op.getResult());
        } else if (!elem_type.isF32()) {

  llvm::SmallVector<int32_t> result_shape_i32;
  auto result_shape = result_ty.getShape();
  for (auto element : result_shape) {
    result_shape_i32.push_back(static_cast<int32_t>(element));
  }
  auto result_shape_attr = builder.getI32TensorAttr(result_shape_i32);
  Value result_shape_tensor =
      builder.create<stablehlo::ConstantOp>(result_shape_attr);
  auto custom_code =

    llvm::SmallVector<Value, 4> new_operands;
    for (auto arg : llvm::enumerate(compose_func_type.getInputs())) {
      if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(arg.value())) {
        auto casted = builder.create<CastOp>(op->getLoc(), tensor_type,
                                             op->getOperand(arg.index()));
        new_operands.push_back(casted);
      } else if (auto list_type =

  auto function_uses =
      SymbolTable::getSymbolUses(target_func, &module_op.getBodyRegion());
  for (auto& function_use : function_uses.value()) {
    auto call_op = function_use.getUser();
    auto caller_func = call_op->getParentOfType<func::FuncOp>();
    if (!caller_func) return failure();

    builder.setInsertionPoint(call_op);
    for (auto [lifted_op, arg_idx] : lifted_op_and_arg_idx) {

  // shaped tensors in the function.
  for (auto arg : func_op.getArguments()) {
    if (!mlir::isa<mlir::TensorType>(arg.getType())) continue;
    auto type = mlir::cast<mlir::TensorType>(arg.getType());
    if (type.hasRank()) {
      max_arg_size_ = std::max(max_arg_size_, GetRankedTensorSize(type));
    }
  }
}

void CostAnalysis::AnalyzeBlock(mlir::Block* block) {
  for (auto& op : *block) {
    EvaluateCost(&op);
  }

      return WalkResult::advance();
    }

    auto backend_config =
        op->getAttrOfType<DictionaryAttr>(kTfBackendConfigAttrName);
    if (!backend_config) {
      op->emitOpError() << "is missing attribute '" << kTfBackendConfigAttrName
                        << "'";
      return WalkResult::interrupt();
    }
    auto called_func = mlir::dyn_cast_or_null<SymbolRefAttr>(

}

// Checks if an operation of the tf_executor dialect can have TPU devices
// propagated through.
bool IsSupportedExecutorOp(Operation& op) {
  auto ops_have_same_device = [](Operation* lhs, Operation* rhs) {
    auto lhs_device_attr = lhs->getAttrOfType<StringAttr>(kDeviceAttr);
    auto rhs_device_attr = rhs->getAttrOfType<StringAttr>(kDeviceAttr);
    return (!lhs_device_attr && !rhs_device_attr) ||

    llvm::SmallVectorImpl<std::string>* var_handle_shared_names) {
  Block& block = function.front();
  auto func_type = function.getFunctionType();

  auto func_arg_types = llvm::to_vector<4>(func_type.getInputs());
  llvm::SmallDenseMap<llvm::StringRef, int> var_arg_index_by_name;
  for (auto var_handle_op :
       llvm::make_early_inc_range(block.getOps<TF::VarHandleOp>())) {

    // (element type of) the expressed type.
    if (auto quantizedType = mlir::dyn_cast<QuantizedType>(spec))
      return quantizedType.isCompatibleExpressedType(expressed);

    if (auto tensorType = mlir::dyn_cast<TensorType>(expressed))
      return spec == tensorType.getElementType();

    if (auto vectorType = mlir::dyn_cast<VectorType>(expressed))
      return spec == vectorType.getElementType();