const Node* identity_input_node = nullptr;
  TF_EXPECT_OK(identity_node->input_node(0, &identity_input_node));
  EXPECT_EQ(identity_input_node->name(), "_arg0");

  const Node* arg_node = node_name_index.at("_arg0");
  EXPECT_EQ(arg_node->op_def().name(), "_Arg");
  EXPECT_EQ(arg_node->attrs().FindByString("T")->type(), DT_FLOAT);

  const Node* retval_node = node_name_index.at("_retval0");

    func.func @main(%arg0: tensor<3x3x10xbf16>, %arg1: tensor<3xi32>) -> tensor<1x?x4xbf16> attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "_arg0,_arg1,_arg2", outputs = "_retval0"}} {
      %cst = "tf.Const"() {value = dense<[1, -1, 4]> : tensor<3xi32>} : () -> tensor<3xi32>

  "tf_device.cluster"() ({

    // CHECK: "tf.Const"() <{value = dense<(1.000000e+00,0.000000e+00)> : tensor<complex<f32>>}> : () -> tensor<complex<f32>>
    // CHECK-NOT: tf.ResourceApplyAdam

  %0:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<1xi32>, tensor<1xf32>) -> (tensor<1xi32>, tensor<1xf32>)
  func.return %0#0, %0#1: tensor<1xi32>, tensor<1xf32>
}

// -----

// CHECK-LABEL: func @stopgradient
func.func @stopgradient(%arg0: tensor<1xi32>) -> tensor<1xi32> {
  // CHECK-NEXT:  return %arg0 : tensor<1xi32>
  %0 = "tf.StopGradient"(%arg0) : (tensor<1xi32>) -> tensor<1xi32>

  %while = "tf.WhileRegion"(%max_size) ({
    // CHECK: ^bb0(%[[BARG0:.*]]: tensor<i32>
    ^bb0(%barg0: tensor<i32>):
     // CHECK: "tf._SomeOp"(%[[BARG0]])
     %pred = "tf._SomeOp"(%barg0) : (tensor<i32>) -> tensor<i1>
    "tf.Yield"(%pred) : (tensor<i1>) -> ()
  }, {
    // CHECK: ^bb0(%[[BARG0:.*]]: tensor<i32>
    ^bb0(%barg0: tensor<i32>):

  }
  func.func @partitioned_called_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
    func.return %arg0 : tensor<?xi32>
  }
  func.func @stateful_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
    func.return %arg0 : tensor<?xi32>
  }
  func.func @tpu_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
    func.return %arg0 : tensor<?xi32>
  }

  // CHECK: return %[[BARG0]], %[[SUB]], %[[BARG2]]
  func.return %arg0, %sub : tensor<!tf_type.resource>, tensor<i32>
}
// CHECK: func @while_cond(%[[CARG0:.*]]: tensor<!tf_type.resource<tensor<5x3xf32>>>, %[[CARG1:.*]]: tensor<i32>, %[[CARG2:.*]]: tensor<!tf_type.resource<tensor<5x3xf32>>>)
func.func @while_cond(%arg0: tensor<!tf_type.resource>, %arg1: tensor<i32>) -> tensor<i32> {

  func.return
}
// CHECK: func @while_body(%[[BARG0:.*]]: tensor<f32>)
func.func @while_body(%arg0: tensor<*x!tf_type.resource<tensor<f32>>>) -> (tensor<*x!tf_type.resource<tensor<f32>>>) {
  // CHECK-NEXT: return %[[BARG0]]
  func.return %arg0 : tensor<*x!tf_type.resource<tensor<f32>>>
}
// CHECK: func @while_cond(%[[CARG0:.*]]: tensor<f32>)
func.func @while_cond(%arg0: tensor<*x!tf_type.resource<tensor<f32>>>) -> tensor<f32> {

func.func @while_shape_invariant_outside_cluster(%arg0: tensor<4xf32>) -> (tensor<*xf32>, tensor<*xf32>) {
  // IN-CLUSTER: shape_invariant
  %0 = "tf.While"(%arg0) {cond = @while_cond, body = @while_body, is_stateless = false, shape_invariant} : (tensor<4xf32>) -> (tensor<*xf32>)

  // IN-CLUSTER: shape_invariant
  %1 = "tf.WhileRegion"(%arg0) ({
  ^cond(%carg0: tensor<*xf32>):

    {
      ^bb0(%carg0: tensor<*xf32>, %carg1: tensor<i32>):
      %ne = "tf.NotEqual"(%carg1, %zero) : (tensor<i32>, tensor<i32>) -> tensor<i1>
      "tf.Yield"(%ne) : (tensor<i1>) -> ()
    },
    {
      // loop body
      ^bb0(%barg0: tensor<*xf32>, %barg1: tensor<i32>):
      %add = "tf.Add"(%barg0, %barg0) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>