}
      }

      // Replace the arg with a tf.Const op in the function body.
      builder.setInsertionPointToStart(&func.getBody().front());
      auto const_op = builder.create<TF::ConstOp>(global_tensor.getLoc(),
                                                  *global_tensor.getValue());
      args_to_erase.set(val.getArgNumber());
      for (auto read_op : read_variable_ops_to_erase) {

  return strings::StrCat(forward_frame_name, "_backprop");
}

// Creates a loop that counts the number of iterations performed by the
// while loop associated with `while_ctx`. The returned output yields the
// iteration count.
Status AddForwardLoopCounter(WhileContext* while_ctx, const Scope& scope,
                             Output* count) {
  // Create while loop:
  //   i = 0
  //   while forward loop predicate is true:

      // Add dequantize.
      builder.setInsertionPointAfter(read_variable_op);
      auto new_read_variable_op =
          builder.create<ReadVariableOp>(read_variable_op.getLoc(), ref_qtype,
                                         read_variable_op.getResourceId());
      auto new_dq_op = builder.create<DequantizeOp>(
          read_variable_op.getLoc(), read_variable_op.getResult().getType(),
          new_read_variable_op.getResult());

    }
    DenseElementsAttr attr =
        DenseElementsAttr::get(type, b.getI32IntegerAttr(index));
    auto constant = b.create<arith::ConstantOp>(op.getLoc(), type, attr);
    op.replaceAllUsesWith(constant.getOperation());
    op.erase();
  });
}

// Creates an instance of the TensorFlow DeviceIndex selector pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDeviceIndexSelectorPass() {

  // available during compilation or compilation only device for on demand
  // execution which may create a recursion if used for constant folding.
  std::string host_cpu = tensorflow::DeviceNameUtils::FullName(
      /*job=*/"localhost", /*replica=*/0, /*task=*/0, /*type=*/"CPU", /*id=*/0);

  absl::StatusOr<OpKernelRunner> runner = OpKernelRunner::Create(
      node_def->get()->op(), node_def->get()->name(), host_cpu, operands.size(),

    for (int i = 0; i < counter_size; ++i) {
      Value word = counter[i];
      Value word_u64 = rewriter.create<CastOp>(loc, word_u64_type, word);
      Value new_word_u64 = rewriter.create<AddV2Op>(loc, word_u64, increment);
      Value new_word = rewriter.create<CastOp>(loc, word_type, new_word_u64);
      pack_args.push_back(new_word);

    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);

    SmallVector<int64_t, 8> line_nums;
    line_nums.resize(lines.size());
    std::iota(line_nums.begin(), line_nums.end(), 0);

    // Create constant ops for keys an values.
    Value key_constant_tensor = rewriter.create<arith::ConstantOp>(
        op.getLoc(),
        DenseStringElementsAttr::get(
            RankedTensorType::get(static_cast<int64_t>(lines.size()),