// CHECK-SAME: op_key = 1
  // CHECK-NOT: __op_key
  %a = "tf.AddV2"(%input0, %const) {__op_key = 1: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
  // CHECK: [[b:%.*]] = tf_mlrt.executeop([[a]])
  // CHECK-SAME: Abs
  // CHECK-SAME: op_key = 2
  // CHECK-NOT: __op_key
  %b = "tf.Abs"(%a) {__op_key = 2: i32 }: (tensor<i32>) -> tensor<i32>
  // CHECK: tf_mlrt.promise [[promise_b]], [[b]]

  %0 = "tf.ReadVariableOp"(%arg1) {__op_key = 0: i32, device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
  %1 = "tf.Shape"(%arg0) {__op_key = 1: i32, device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
  %2 = "tf.Shape"(%0) {__op_key = 2: i32, device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>

// CHECK-LABEL: func @main
// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 0 : i32}

// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 1 : i32}

// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 2 : i32}

// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 3 : i32}

// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 4 : i32}

// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 5 : i32}

// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 6 : i32}

func.func @while_body_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>) {
  %0 = "tf.Const"() {__op_key = 0: i32, device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
  %1 = "tf.Add"(%z, %0) {__op_key = 1: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %cond, %x, %y, %1 : tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>
}

#define TENSORFLOW_COMPILER_MLIR_TFRT_CONSTANTS_H_

namespace tensorflow {
namespace tfrt_compiler {

// Use __ prefix to indicate this is internal attribute.
inline constexpr char kOpKeyAttrName[] = "__op_key";

}  // namespace tfrt_compiler
}  // namespace tensorflow

  EXPECT_EQ(expected_op_cost_map.size(), 1);
  unsigned int seed = 23579;
  for (auto& [op_key, cost] : expected_op_cost_map) {
    cost = rand_r(&seed) % 1000;
  }
  tensorflow::tfrt_stub::CostRecorder cost_recorder;
  for (const auto& [op_key, cost] : expected_op_cost_map) {
    cost_recorder.RecordCost(op_key, cost);
  }

  // Update the TFRT MLIR with the cost recorder.

def CreateOp: FallbackAsync_Op<"createop", [CoreRT_TypedAttributeTrait]> {
  let summary = "The Fallback CreateOp";

  let description = [{
    The CreateOp creates the tensorflow::OpKernel in the fallback context.

    op_key: The unique key to identify this op.
    num_args: The number of inputs.
    device: The tensorflow device. eg. "/CPU:0"
    op_attrs: The tensorflow attributes excluding the func attrs.
    op_func_attrs: The func attrs.

                              mlir::OpAsmParser::Delimiter::Paren))
    return mlir::failure();

  if (options.has_key &&
      (parser.parseKeyword("key") || parser.parseLParen() ||
       parser.parseAttribute(op_key, "op_key", result.attributes) ||
       parser.parseRParen()))
    return mlir::failure();

  if (options.has_cost &&
      (parser.parseKeyword("cost") || parser.parseLParen() ||

  let summary = "The Fallback CreateOp";

  let description = [{
    The CreateOp creates the tensorflow::OpKernel in the fallback context.
  }];

  let arguments = (ins
    StrAttr:$node_def,
    I32Attr:$op_key
  );

  let assemblyFormat = "attr-dict";
}

def ConstOp: TensorflowMlrt_Op<"constop", []> {
  let summary = "The tf_mlrt ConstOp";

  let description = [{

  // CHECK-NEXT: [[t3:%.*]] = tf_mlrt.await [[f2]]
  // CHECK-NEXT: return [[t2]], [[t3]]
  %t0 = tf_mlrt.await %f0
  %t1 = tf_mlrt.await %f1
  %t2 = tf_mlrt.executeop(%t0, %t1) {node_def = "AddV2", op_key = 0 : i32} : (!tf_mlrt.tensor, !tf_mlrt.tensor) -> (!tf_mlrt.tensor)
  %t3 = tf_mlrt.await %f2
  func.return %t2, %t3 : !tf_mlrt.tensor, !tf_mlrt.tensor
}

// -----

// CHECK-LABEL: @main