//          (tensor<?x3x4x2xf32>) -> tensor<i32>
// %5 = stablehlo.reshape %4 :
//          (tensor<i32>) -> tensor<1xi32>
// %6 = stablehlo.concatenate %5, %0, %1, %2, dim = 0 :
//          (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>)
//            -> tensor<4xi32>
// %7 = stablehlo.dynamic_broadcast_in_dim %arg2, %6
// %8 = stablehlo.add %3, %7
// ```
//

    %0 = "tfl.pseudo_const"() {value = dense<0.962260901> : tensor<1xf32>} : () -> tensor<1xf32>
    %1 = func.call @func_0_GPU_FLOAT(%arg0, %0) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1x200x200x200xf32>, tensor<1xf32>) -> tensor<1x200x200x200xf32>
    %2 = "tfl.pseudo_const"() {value = dense<0.895973444> : tensor<1xf32>} : () -> tensor<1xf32>

    %3 = "tf.SomeComputation"(%2) : (tensor<*xi32>) -> (tensor<*xi32>)
    "tf.AssignVariableOp"(%0, %3) {dtype = i32} : (tensor<*x!tf_type.resource<tensor<4xi32>>>, tensor<*xi32>) -> ()
    tf_device.return %3 : tensor<*xi32>
  }) {cluster_attr = "cluster_attr"} : () -> tensor<*xi32>

  // CHECK: return %[[CLUSTER_RES]]#0
  // CHECK-SAME: tensor<*xi32>

    ```mlir
      %2 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
      %3 = "tf.B"(%2) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>
      %4 = "tf.C"(%2, %3) {device = "tpu0"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
      %5 = "tf.D"(%4) : (tensor<?xi32>) -> tensor<?xi32>
    ```

    After the pass, we will have:

    ```mlir
      %0 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>

// CHECK-DAG: %[[BROADCAST_DIM:.+]] = arith.constant dense<{{\[3, 2, 1, 1\]}}> : tensor<4xi32>
// CHECK-DAG: %[[EXPAND_DIM1:.+]] = arith.constant dense<3> : tensor<1xi32>
// CHECK-DAG: %[[EXPAND_DIM0:.+]] = arith.constant dense<2> : tensor<1xi32>
// CHECK: %[[EXPAND0:.+]] = "tfl.expand_dims"(%[[ARG0]], %[[EXPAND_DIM0]]) : (tensor<1x2x!quant.uniform<i8:f32, 2.000000e+00:3>>, tensor<1xi32>) -> tensor<1x2x1x!quant.uniform<i8:f32, 2.000000e+00:3>>

    %cst_1 = "tf.Const"() {value = dense<[2, 1024]> : tensor<2xi32>} : () -> tensor<2xi32>
    // Original identity preserved.
    %cst_identity = "tf.Identity"(%cst_0) {device = ""} : (tensor<1024x2xf32>) -> tensor<1024x2xf32>
    %0 = "tf.Reshape"(%cst_identity, %cst_1) : (tensor<1024x2xf32>, tensor<2xi32>) -> tensor<2x1024xf32>

```mlir
  %2 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
  %3 = "tf.B"(%2) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>
  %4 = "tf.C"(%2, %3) {device = "tpu0"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
  %5 = "tf.D"(%4) : (tensor<?xi32>) -> tensor<?xi32>
```

After the pass, we will have:

```mlir
  %0 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
  %1 = "tf_device.launch"() ( {

    %s0 = "tf.Const"() {value = dense<[501, 1, 32, 1280]> : tensor<4xi32>} : () -> tensor<4xi32>
    %s1 = "tf.Const"() {value = dense<[  1, 1,  1, 1280]> : tensor<4xi32>} : () -> tensor<4xi32>

    // CHECK: "tf.BroadcastGradientArgs"
    // CHECK-NOT: _xla_outside_compilation
    %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) {} : (tensor<4xi32>, tensor<4xi32>) -> (tensor<1xi32>, tensor<3xi32>)
    tf_device.return

  // expected-remark@above {{Sinks: {}}}
}

// -----

func.func @add(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  // expected-remark@above {{ID: 2}}
  %sum = "tf.Add"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  // expected-remark@above {{ID: 0}}
  func.return %sum : tensor<1xf32>
  // expected-remark@above {{ID: 1}}
  // expected-remark@above {{Sinks: {}}}
}

  %cst_1 = "tf.Const"() {value = dense<[2, 0, 1]> : tensor<3xi32>} : () -> tensor<3xi32>
  %cst_2 = "tf.Const"() {value = dense<[512, 64]> : tensor<2xi32>} : () -> tensor<2xi32>
  %0 = "tf.ExpandDims"(%arg0, %cst) {device = ""} : (tensor<8x64xf32>, tensor<i32>) -> tensor<8x64x1xf32>
  %1 = "tf.AddV2"(%0, %arg1) {device = ""} : (tensor<8x64x1xf32>, tensor<8x64x64xf32>) -> tensor<8x64x64xf32>