func.func @add_v2_uint32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
  func.return %0 : tensor<ui32>

  // CHECK-LABEL: add_v2_uint32
  // CHECK:  %[[CAST:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<ui32>) -> tensor<i32>
  // CHECK:  %[[CAST1:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<ui32>) -> tensor<i32>

  %outputs_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
  %outputs_2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
  %outputs_4 = "tf.Identity"(%arg0) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
  %outputs_6 = "tf.AddV2"(%outputs_4, %outputs_2) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
  %outputs_8 = "tf.Identity"(%arg1) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>

  %3:2 = "tf.TPUReplicatedOutput"(%2#0) : (tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>)
  %5:2 = "tf.TPUPartitionedOutputV2"(%3#0) {_XlaSharding = "", partition_dims = []} : (tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>)
  %7:2 = "tf.TPUPartitionedOutputV2"(%3#1) {_XlaSharding = "", partition_dims = []} : (tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>)

    }) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>, tensor<?xi32>, tensor<?xi32>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?xi32>>>, tensor<!tf_type.variant<tensor<?xi32>>>) -> (tensor<i32>, tensor<i32>, tensor<?xi32>, tensor<?xi32>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>,...

    // CHECK: [[TENSOR:%.+]] = "tf.TensorScatterAdd"([[READ]], [[EXPAND]], [[UPDATE]]) : (tensor<*xi32>, tensor<?x1xi32>, tensor<?x?x?xi32>) -> tensor<*xi32>
    // CHECK: "tf.AssignVariableOp"([[VAR]], [[TENSOR]])
    "tf.ResourceScatterAdd"(%resource, %indices, %updates) : (tensor<*x!tf_type.resource<tensor<*xi32>>>, tensor<?xi32>, tensor<?x?x?xi32>) -> ()

    tf_device.return

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

  // CHECK: return %[[CLUSTER_RES]]
  return %1 : 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>
  %1 = "tf_device.launch"() ( {

    %31 = "tf.Mul"(%30, %cst_7) : (tensor<1xi32>, tensor<i32>) -> tensor<1xi32>
    %32 = "tf.Add"(%24, %27) : (tensor<2x?x1xi32>, tensor<2x1x?xi32>) -> tensor<2x?x?xi32>
    %33 = "tf.Sub"(%32, %31) : (tensor<2x?x?xi32>, tensor<1xi32>) -> tensor<2x?x?xi32>
    %34 = "tf.Sub"(%21, %33) : (tensor<2x?x?xi32>, tensor<2x?x?xi32>) -> tensor<2x?x?xi32>
    %35 = "tf.Cast"(%34) {Truncate = false} : (tensor<2x?x?xi32>) -> tensor<2x?x?xf32>

  // CHECK-DAG: %[[RANGE_M:.*]] = "tf.Range"(%[[ZERO]], %[[M]], %[[ONE]]) : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
  // CHECK-DAG: %[[RANGE_N:.*]] = "tf.Range"(%[[ZERO]], %[[N]], %[[ONE]]) : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
  // CHECK-DAG: %[[EXPAND_DIMS:.*]] = "tf.ExpandDims"(%[[RANGE_M]], %[[NEG_ONE]]) : (tensor<?xi32>, tensor<i32>) -> tensor<?x1xi32>

      name = "avgpool", filter_height = 3 : i32, filter_width = 6 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 3 : i32, stride_w = 1 : i32
    } : (tensor<1x224x224x3xf32>) -> tensor<1x224x224x3xf32>
  %6 = "tfl.conv_2d"(%5, %4, %cst) {
      dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32