auto builder = OpBuilder::atBlockTerminator(&main_graph_op.GetBody());
  // Create an IslandOp that will wrap the IdentityOp. Add a control dependency
  // for the newly copied save function.
  auto wrapper_island_op = builder.create<IslandOp>(
      name_loc, TypeRange{main_file_prefix_arg.getType()},
      tf_executor::ControlType::get(&ctx), ValueRange(control_inputs));

  const std::string dump_file_path =
      tsl::io::JoinPath(*dump_dir, dump_file_name);
  TF_ASSIGN_OR_RETURN(std::unique_ptr<llvm::raw_ostream> file,
                      WritableFileWrapper::Create(dump_file_path));

  LOG(INFO) << "IR dump file created: " << dump_file_path;
  return file;
}

class PrinterConfig : public mlir::PassManager::IRPrinterConfig {
 public:
  explicit PrinterConfig(

    auto begin_op =
        rewriter.create<arith::ConstantOp>(op->getLoc(), type, begin_attr);
    auto end_attr = DenseElementsAttr::get<int32_t>(type, padded_end);
    auto end_op =
        rewriter.create<arith::ConstantOp>(op->getLoc(), type, end_attr);
    auto stride_attr = DenseElementsAttr::get<int32_t>(type, padded_stride);
    auto stride_op =
        rewriter.create<arith::ConstantOp>(op->getLoc(), type, stride_attr);

    // If the op is rewritten, FusedLoc can be created.
    for (Location child_loc : mlir::cast<FusedLoc>(loc).getLocations()) {
      std::optional<QuantizationUnitLoc::QuantizationUnit> found_unit =
          FindQuantizationUnitFromLoc(child_loc);
      if (found_unit.has_value()) return found_unit;
    }
  } else if (isa<CallSiteLoc>(loc)) {
    // If the graph is inlined, CallSiteLoc can be created.
    return FindQuantizationUnitFromLoc(

  OpBuilder builder(merged_island.insert_point);
  auto new_island = builder.create<IslandOp>(
      merged_island.insert_point->getLoc(), result_types, operands);
  new_island.getBody().push_back(new Block);
  return new_island;
}

// Creates respective YieldOp for the new merged island.
YieldOp CreateNewIslandYieldOp(IslandOp new_island,

    const Type expressed_type = value.getType();
    // The value needs to be requantized. A Quantize op will be created to use
    // it as the operand and replace its uses.
    const Type new_type = state.params.castFromExpressedType(expressed_type);
    if (!new_type) return;
    auto requantize_op =
        builder_.create<quantfork::QuantizeCastOp>(loc, new_type, value);
    value.replaceAllUsesWith(requantize_op);

      report_file_path));

  const LogicalResult run_result = pm.run(*module_op);
  ASSERT_TRUE(succeeded(run_result));

  // The report file is not created because `QuantizeCompositeFunctionsPass` was
  // not run.
  EXPECT_THAT(ReadFileToString(report_file_path),
              StatusIs(absl::StatusCode::kNotFound));
}

TEST_F(SaveQuantizationReportInstrumentationTest,

  // Guarantee all functions have one use, which enables more exact shape
  // inference.
  pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
  // Run shape inference so that tf_executor/tf_device ops created later will
  // likely to inherit more concrete types.
  pm.addPass(mlir::TF::CreateTFShapeInferencePass());
  pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUPartitionedOpConversionPass());
  pm.addNestedPass<FuncOp>(

          "and will trigger an error in the future. Either don't modify "
          "nodes after running them or create a new session.");
    }
  }
}

namespace {

// Helper method that creates a shape handle for a shape described by dims.
tensorflow::shape_inference::ShapeHandle ShapeHandleFromDims(
    tensorflow::shape_inference::InferenceContext* ic, int num_dims,

    return tokens[0];
  } else {
    return builder.create<AfterAllOp>(loc, original_token.getType(), tokens)
        .getResult();
  }
}

// Replaces `tf._XlaHostComputeMlir` with individual `mhlo.send` and `mhlo.recv`
// ops per operand and result. Unique Channel IDs are assigned per transfer.
// Sink tokens are created across all `mhlo.send` ops first and then by
// all `mhlo.recv` ops.