#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"

namespace tensorflow {
namespace tf_mlrt {

REGISTER_OP("TestAsyncIdentity")
    .Input("in: T")
    .Output("out: T")
    .Attr(
        "T: {bfloat16, half, float, double, uint8, int8, int16, uint32, int32, "
        "int64, complex64, complex128}")

// Test lowering AsyncOpKernel

// CHECK-LABEL: func @main
func.func @main(%x: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
  // CHECK: [[y_future:%.*]] = tf_mlrt.async_executeop
  %y = "tf.TestAsyncIdentity"(%x) {__op_key = 0: i32, T = i32} : (tensor<i32>) -> tensor<i32>
  // CHECK: [[z:%.*]] = tf_mlrt.executeop
  %z = "tf.Identity"(%x) {__op_key = 1: i32}: (tensor<i32>) -> tensor<i32>