%0 = "tf.StatefulPartitionedCall"(%arg0) <{
      config = "", config_proto = "", executor_type = "", f = @some_func
    }> {
      _collective_manager_ids = [], device = ""
    } : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    func.return %0: tensor<1x2x2x3xf32>
  }

  func.func private @some_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> attributes {tf._noinline = true} {
    return %arg0 : tensor<1x2x2x3xf32>

  ) attributes { tf_saved_model.exported_names = ["foo.some_func"] } {
    func.return
  }

}

// -----

// expected-error@+1 {{'tf_saved_model.exported_names' must be on an op whose immediate parent has attribute 'tf_saved_model.semantics'}}
func.func @f() attributes { tf_saved_model.exported_names = ["foo.some_func"] } {
  func.return
}

// -----

  // CHECK:  "tf._TfrtGetResource"()
  %cond = "tf.Const"() {device = "/CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
  // CHECK: tf.If
  %x = "tf.If"(%cond, %handle) {then_branch = @some_func, else_branch = @some_func, is_stateless = false} : (tensor<i1>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  %r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %r : tensor<i32>

}

}

// -----

module attributes {tf_saved_model.semantics} {

// Test sinks in var handle op that used by control flow ops.

// CHECK-LABEL: func private @some_func
func.func private @some_func(
    %arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
  // CHECK: tf.VarHandleOp
  // CHECK: tf.ReadVariableOp

    %1 = stablehlo.uniform_quantize %arg0 : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
    // CHECK: call @some_func
    // CHECK-SAME: -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>

  // CHECK: %[[ADD:.*]] = "tf.StatefulPartitionedCall"(%[[INPUT]]#1, %[[INPUT]]#2)
  // CHECK-SAME: f = @some_func
  // CHECK: "tf.PartitionedCall"(%[[ADD]], %[[INPUT]]#0)
  // CHECK-SAME: f = @some_other_func
  // CHECK: return

  func.func private @some_func(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<3x3xf32> attributes {tf._noinline = true} {

TEST_P(UnifiedCAPI, TF_AbstractOpSetOpTypeAfterFinishingOpBuildingRaises) {
  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);
  TF_ExecutionContext* graph_ctx = TF_CreateFunction("some_func", status.get());
  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());

  // Add a placeholder to the graph.
  auto* placeholder_op = TF_NewAbstractOp(graph_ctx);

  %0 = "tf.MyMapAndBatchDataset"(%input, %other1, %other2)
    {batch_size=1000 : i64, num_parallel_calls = 8 : i64, drop_remainder = 0 : i1,
     func = @"__some_func", output_types = [f32], output_shapes = [#tf_type.shape<>], preserve_cardinality = true}
    : (tensor<*x!tf_type.variant>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.variant>
  func.return %0 : tensor<*x!tf_type.variant>

  FunctionDefLibrary library;

  FunctionDef sort_func = FunctionDefHelper::Create(
      "SortFunc",
      /*in_def=*/{"x:float"},
      /*out_def=*/{"res:float"},
      /*attr_def=*/{},
      /*node_def=*/{{{"t0"}, "XlaSort", {"x"}, {{"T", DT_FLOAT}}}},
      /*ret_def*/ {{"res", "t0:output"}});

  *library.add_function() = sort_func;

  Output in = ops::Placeholder(root, DT_FLOAT);