func.func @variant_input_output_already_marked() {
  // CHECK: "tf.opA"
  // CHECK-SAME: _xla_outside_compilation = "0"
  // CHECK: "tf.opB"
  // CHECK-SAME: _xla_outside_compilation = "0"
  "tf_device.cluster"() ({
    %1= "tf.opA"() {_xla_outside_compilation = "0"} : () -> tensor<!tf_type.variant<tensor<f32>>>
    "tf.opB"(%1) {_xla_outside_compilation = "0"} : (tensor<!tf_type.variant<tensor<f32>>>) -> ()
    tf_device.return

  // CHECK-NEXT: "tf.Cast"
  // CHECK-SAME: _xla_outside_compilation = ""
  // CHECK-NEXT: "tf.Identity"
  // CHECK-SAME: _xla_outside_compilation = ""
  // CHECK-NEXT: "tf.B"
  "tf_device.cluster"() ({
    %1 = "tf.Cast"(%arg0) : (tensor<?xi32>) -> (tensor<?xi32>)
    %2 = "tf.Identity"(%1) : (tensor<?xi32>) -> (tensor<?xi32>)
    "tf.B"(%2) {_xla_outside_compilation = "cluster1"} : (tensor<?xi32>) -> ()
    "tf.C"() : () -> ()

    %0 = "tf.Shape"(%arg0) {_tpu_replicate = "cluster_Fn", _xla_outside_compilation = "1", device = ""} : (tensor<*x!tf_type.string>) -> tensor<*xi32>
    %1 = "tf.Shape"(%arg1) {_tpu_replicate = "cluster_Fn", _xla_outside_compilation = "0", device = ""} : (tensor<*x!tf_type.string>) -> tensor<*xi32>

    // CHECK-NEXT:   %[[A_OUT:.*]] = "tf.A"
    // CHECK-NOT:    _xla_outside_compilation
    // CHECK-NEXT:   %[[B_OUT:.*]] = "tf.B"(%[[A_OUT]])
    // CHECK-NOT:    _xla_outside_compilation
    // CHECK-NEXT:   %[[C_OUT:.*]] = "tf.C"(%[[B_OUT]], %arg0)
    // CHECK-NOT:    _xla_outside_compilation
    // CHECK-NEXT:   tf_device.return %[[C_OUT]]
    //
    // CHECK:      "tf_device.cluster"

tensor<1120x?xi32>, tensor<2xi32>) {\0A    %0 = \22tf.Reshape\22(%arg0, %arg1) {_xla_outside_compilation = \220\22} : (tensor<3360x?xi32>, tensor<3xi32>) -> tensor<3x1120x?xi32> loc(#loc9)\0A    %1:3 = \22tf.Split\22(%arg2, %0) {_xla_outside_compilation = \220\22} : (tensor<i32>, tensor<3x1120x?xi32>) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1x1120x?xi32>) loc(#loc10)\0A    %2 = \22tf.Reshape\22(%1#0, %arg3) {_xla_outside_compilation = \220\22} : (tensor<1x1120x?xi32>, tensor<2xi32>) -> tensor<1120x?xi32>...

     // CHECK-NEXT: "tf.D"
     // CHECK-NOT: _xla_outside_compilation
     // CHECK: "tf_device.cluster"
     // CHECK-NEXT: "tf.A"
     // CHECK-NEXT: "tf.E"
     // CHECK: device_assignment =  [], num_cores_per_replica = 1 : i64, topology =  ""
    "tf_device.cluster"() ({
      "tf.A"() : () -> ()
      "tf.B"() {_xla_outside_compilation = "cluster1"} : () -> ()
      "tf.C"() {_xla_outside_compilation = "cluster1"} : () -> ()

    %3 = "tf.All"(%0, %2) {_tpu_replicate = "cluster", _xla_outside_compilation = "0", device = "", keep_dims = false} : (tensor<*xi1>, tensor<*xi32>) -> tensor<*xi1>

  // CHECK: %[[CONST_0:.*]] = "tf.Const"()
  %0 = "tf.Const"() {value = dense<[]> : tensor<0xf32>} : () -> tensor<0xf32>

  // CHECK: %[[CONST_MODE:.*]] = "tf.Const"() <{value = dense<"inference"> : tensor<!tf_type.string>}> {_xla_outside_compilation = "0"} : () -> tensor<!tf_type.string>
  // CHECK: "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_3]], %[[ARG_4]], %[[ARG_5]], %[[CONST_0]], %[[CONST_0]], %[[CONST_0]], %[[CONST_MODE]])

// -----

func.func @testTFDialect(%arg0: tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> {
   // expected-error@below {{op has outside compilation attribute _xla_outside_compilation which is not allowed after clustering}}
  %0 = "tf.Identity"(%arg0) {_xla_outside_compilation = "cluster1"}: (tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string>
  func.return %0 : tensor<4x2x!tf_type.string>

    For example, the following tf_device.cluster with an op marked for `xla_outside_compilation`:

    ```mlir
    func @outside_compilation() -> tensor<f32> {
      %0 = "tf_device.cluster"() ( {
        %1 = "tf.Const"() {_xla_outside_compilation = "0", value = dense<1.0> : tensor<f32>} : () -> (tensor<f32>)
        %2 = "tf.Identity"(%1) {_xla_outside_compilation = "0"} : (tensor<f32>) -> (tensor<f32>)