b.setInsertionPoint(op);
    if (auto dq = llvm::dyn_cast<DequantizeOp>(op)) {
      auto dcast = b.create<quantfork::DequantizeCastOp>(
          dq.getLoc(), dq.getOutput().getType(), dq.getInput());
      dq.getOutput().replaceAllUsesWith(dcast);
      dq.erase();
    } else if (auto q = llvm::dyn_cast<QuantizeOp>(op)) {
      auto qcast = b.create<quantfork::QuantizeCastOp>(
          q.getLoc(), q.getOutput().getType(), q.getInput());

  auto runtime =
      tensorflow::tfrt_stub::Runtime::Create(/*num_inter_op_threads=*/1);
  tfrt_stub::GraphExecutionOptions options(runtime.get());

  options.compile_options.device_target = TfrtDeviceInfraTarget::kGpu;

  TF_ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<tfrt_stub::FallbackState> fallback_state,
      tfrt_stub::FallbackState::Create(SessionOptions(), FunctionDefLibrary()));

    f.setName(StringAttr::get(f->getContext(), f.getName() + "_internal"));

    for (StringRef name : llvm::reverse(exported_names)) {
      // Create a "trampoline" function with the given name. So given
      //   func bar(...) {exported_names = ["foo"]}
      // we create
      //   func foo(...) {
      //     return bar(...)
      //   }
      func::FuncOp other = f.cloneWithoutRegions();
      other.setName(name);

                    llvm::StringRef device) {
  builder->setInsertionPoint(op);
  auto launch = builder->create<tf_device::LaunchOp>(
      loc, builder->getStringAttr(device), op->getResultTypes());
  launch.getBody().push_back(new Block);
  op->replaceAllUsesWith(launch);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<tf_device::ReturnOp>(loc, op->getResults());

  // Move op inside cluster.

    // the resource, and we use it to just look up.
    // See core/kernels/data/iterator_ops.cc.)
    resource_op = builder.create<TF::IteratorOp>(
        op->getLoc(), resource.getType(), shared_name, container,
        it.getOutputTypes(), it.getOutputShapes());
  } else {
    resource_op = builder.create<TF::VarHandleOp>(
        op->getLoc(), resource.getType(), container, shared_name);
  }

    builder.setInsertionPoint(op);
    for (auto& op_operand : op->getOpOperands()) {
      if (IsLargeFloatType(op_operand.get().getType())) {
        op_operand.set(builder.create<TF::CastOp>(
            op->getLoc(), ToBfloat16Type(op_operand.get().getType()),
            op_operand.get()));
      }
    }
    builder.setInsertionPointAfter(op);

  builder.setInsertionPoint(composite);
  if (IsKVCacheCompositeOp(composite)) {
    return builder.create<TFL::CustomOp>(
        composite->getLoc(), composite->getResultTypes(),
        composite->getOperands().slice(2, 3), composite.getName(),
        CustomOption(&builder, custom_option_buffer));
  }
  return builder.create<TFL::CustomOp>(
      composite->getLoc(), composite->getResultTypes(),

namespace mlir {
namespace odml {

LogicalResult ConvertCustomCallOp::matchAndRewrite(
    mhlo::CustomCallOp mhlo_custom_call, OpAdaptor adaptor,
    ConversionPatternRewriter& rewriter) const {
  auto tfl_custom = rewriter.create<TFL::CustomOp>(
      mhlo_custom_call.getLoc(), mhlo_custom_call.getResultTypes(),
      mhlo_custom_call.getInputs());
  tfl_custom.setCustomCodeAttr(
      rewriter.getStringAttr(mhlo_custom_call.getCallTargetName()));

    const llvm::SmallVector<int64_t>& expected_packed_shape,
    const llvm::SmallVector<int8_t>& expected_packed_values) {
  MLIRContext context;
  OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
  OpBuilder builder(&module->getBodyRegion());
  context.loadDialect<TF::TensorFlowDialect>();

  Value value = CreateConstValue<int8_t>(builder, module->getLoc(),

  mlir::MLIRContext context(registry);
  auto module =
      mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
  ASSERT_TRUE(module);

  // Create a cost recorder with fake cost records.
  auto expected_op_cost_map = GetOpCostMap(module.get());
  EXPECT_EQ(expected_op_cost_map.size(), 1);
  unsigned int seed = 23579;
  for (auto& [op_key, cost] : expected_op_cost_map) {