}

// CHECK-LABEL: func @doubly_nested_embedding_op
func.func @doubly_nested_embedding_op(%arg0: tensor<i1>, %arg1: tensor<i1>, %arg2: tensor<512x256xf32>) -> (tensor<512x256xf32>) {
  %2 = "tf.IfRegion"(%arg0) ({
    %1 = "tf.IfRegion"(%arg1) ({
      // CHECK: "tf.XlaRecvTPUEmbeddingDeduplicationData"
      // CHECK: "tf.XlaRecvTPUEmbeddingActivations"
      // CHECK-NOT: tf.RecvTPUEmbeddingActivations

    // `tf_then_or_else_region` is the `then` or `else` region of the
    // `tf.IfRegion` op and `scf_then_or_else_region` is the `then` or `else`
    // region of the new `scf.if` op. Further, `tf_if_region_return_type` is the
    // list of return types of the `tf.IfRegion` op.
    auto createScfThenOrElse = [](Region& tf_then_or_else_region,
                                  Region& scf_then_or_else_region,

// RUN: tf-opt -convert-tf-control-flow-to-scf %s | FileCheck %s

// `tf.IfRegion` which returns values gets converted to `scf.if`.
func.func private @test_if_then1(tensor<4xf32>) -> tensor<4xf32>
func.func private @test_if_else1(tensor<4xf32>) -> tensor<4xf32>
// CHECK-LABEL: func @test_supported_lowering_of_tf_if_region1
// CHECK-SAME: (%[[ARG0:.*]]: tensor<i1>, %[[ARG1:.*]]: tensor<4xf32>)

  func.return
}

// CHECK-LABEL: @IfOpTokenAttrs
func.func @IfOpTokenAttrs(%arg0: tensor<i1>, %arg1: tensor<f32>) -> (tensor<f32>) {
  // CHECK: tf.IfRegion
  %0 = "tf.IfRegion"(%arg0) ({
      // CHECK: tf.XlaRecvFromHost
      // CHECK-SAME-DAG: _xla_original_oc_node_name =
      // CHECK-SAME-DAG: _xla_token_input_nodes = ["_xla_token_arg_node"]

    // CHECK:          %[[GET_NEXT:[0-9]*]] = "tf.IteratorGetNextAsOptional"(%[[ANON_ITER]])
    // CHECK-NEXT:     %[[HAS_VALUE:[0-9]*]] = "tf.OptionalHasValue"(%[[GET_NEXT]])
    // CHECK-NEXT:     %[[IF:.*]] = "tf.IfRegion"(%[[HAS_VALUE]])
    // CHECK-NEXT:       %[[GET_VALUE:[0-9]*]] = "tf.OptionalGetValue"(%[[GET_NEXT]])
    // CHECK-NEXT:       %[[FUNC_CALL:[0-9]*]] = func.call @__reduce_func_1(%[[ARG_5]], %[[GET_VALUE]])

  // CHECK-NOT: _replication_info =
  // CHECK-NOT: _xla_compile_device_type
  // CHECK-NOT: device =
  %1 = "tf_device.cluster"() ({
    %2 = "tf.Add"(%arg1, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
    %3 = "tf.IfRegion"(%arg0) ({
        %4 = "tf.Mul" (%arg1, %2) {device = "y"}: (tensor<f32>, tensor<f32>) -> tensor<f32>
        "tf.Yield"(%4) : (tensor<f32>) -> ()
      }, {
        // CHECK: device = "/device:TPU_REPLICATED_CORE:0"

                                   llvm::ArrayRef<func::FuncOp> functions,
                                   const BacktrackAnalysis& backtrack_analysis);

  // Analyzes tf.CaseRegion/tf.IfRegion ops to compute resource IDs.
  void AnalyzeRegionCaseOrIfOp(Operation* case_or_if_op,
                               const BacktrackAnalysis& backtrack_analysis);