%add2, %add2_control = tf_executor.island(%add1_control) wraps "tf.Add"(%add1, %arg1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
    %print, %print_control = tf_executor.island wraps "tf.Print"(%add2) {message = "add2 result"} : (tensor<*xi32>) -> tensor<*xi32>
    tf_executor.fetch %add1, %add2, %add2_control : tensor<*xi32>, tensor<*xi32>, !tf_executor.control
  }

  %0 = "tf.Identity"(%cst) {device = ""} : (tensor<1xf32>) -> tensor<1xf32>
  %1 = "tf.AddV2"(%0, %cst_1) {device = ""} : (tensor<1xf32>, tensor<f32>) -> tensor<1xf32>
  %2 = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {strides = [1, 2, 2, 1], padding="SAME", dilations=[1, 1, 1, 1]}: (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32>
  %3 = "tf.AddV2"(%2, %1): (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32>

      *literal2, xla::LiteralUtil::CreateR2<int32_t>({{3, 4}})));
}

TEST_F(PjRtExecutionUtilTest, PopulateCtxOutputs) {
  XlaOpRegistry::RegisterCompilationKernels();
  TF_EXPECT_OK(NodeDefBuilder("AddV2", "AddV2")
                   .Input(FakeInput(DT_INT32))
                   .Input(FakeInput(DT_INT32))
                   .Attr("T", DT_INT32)
                   .Device("/job:localhost/replica:0/task:0/device:XLA_CPU:0")

    %3 = "tf.LookupTableSizeV2"(%1) {device = ""} : (tensor<!tf_type.resource>) -> tensor<i64>
    %4 = "tf.AddV2"(%2, %3) {device = ""} : (tensor<i64>, tensor<i64>) -> tensor<i64>
    %5 = "tf.LookupTableFindV2"(%0, %arg0, %cst) {device = ""} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
    %6 = "tf.AddV2"(%5, %4) {device = ""} : (tensor<*xi64>, tensor<i64>) -> tensor<*xi64>
    return %6 : tensor<*xi64>
  }

    %island1:3 = tf_executor.island {
      %add1 = "tf.Add"(%arg0, %arg1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
      %add2 = "tf.Add"(%add1, %arg1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
      %res = "tf.Print"(%add2) { message = "add result" } : (tensor<*xi32>) -> (tensor<*xi32>)
      tf_executor.yield %add1, %add2 : tensor<*xi32>, tensor<*xi32>
    }
    tf_executor.fetch %island1#0, %island1#1 : tensor<*xi32>, tensor<*xi32>
  }

  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  return %0: tensor<i32>
}

// CHECK-LABEL: func.func private @"random/while_body/TfMlrtAsyncWhileBody"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
func.func private @"random/while_body/TfMlrtAsyncWhileBody"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {

    %0 = "tf.MatMul"(%arg0, %cst_0) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    // AddV2 not in quantizable op list.
    %1 = "tf.AddV2"(%arg1, %cst_0) {device = ""} : (tensor<2x1024xf32>, tensor<2x1024xf32>) -> tensor<2x1024xf32>
    func.return %0, %1 : tensor<*xf32>, tensor<2x1024xf32>
  }

// CHECK-LABEL: func @matmul_multiuses

    %0 = "tf.AddV2"(%arg0, %cst) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    %1 = "tf.Identity"(%0) {device = ""} : (tensor<i32>) -> tensor<i32>
    %2 = "tf.Identity"(%arg1) {device = ""} : (tensor<i32>) -> tensor<i32>
    %4 = "tf.AddV2"(%arg2, %cst_0) {device = ""} : (tensor<1x5x5x1024xf32>, tensor<1x5x5x1024xf32>) -> tensor<1x5x5x1024xf32>

    %30 = "tf.CrossReplicaSum"(%29, %3) : (tensor<f32>, tensor<1x2xi32>) -> tensor<f32>
    %31 = "tf.AddV2"(%arg5, %23) : (tensor<f32>, tensor<f32>) -> tensor<f32>
    %32 = "tf.CrossReplicaSum"(%1, %3) : (tensor<f32>, tensor<1x2xi32>) -> tensor<f32>
    %33 = "tf.AddV2"(%arg6, %32) : (tensor<f32>, tensor<f32>) -> tensor<f32>
    %34 = "tf.AddV2"(%arg7, %30) : (tensor<f32>, tensor<f32>) -> tensor<f32>

    %val = "tf.ReadVariableOp"(%handle) : (tensor<!tf_type.resource<tensor<0xf32>>>) -> tensor<0xf32>
    %0 = "tf.AddV2"(%val, %val) : (tensor<0xf32>, tensor<0xf32>) -> tensor<0xf32>
    "tf.AssignVariableOp"(%handle, %0) : (tensor<!tf_type.resource<tensor<0xf32>>>, tensor<0xf32>) -> ()
    // CHECK: "tf.VarHandleOp"
    func.return
  }
}