func.func @does_reordering_for_nested_compiles() -> () {
   // CHECK: _TPUCompileMlir
   // CHECK-SAME: Z
   // CHECK: tf_device.launch
   // CHECK-NEXT: _TPUCompileMlir
   // CHECK-SAME: X
   // CHECK: tf_device.launch
   // CHECK-NEXT: _TPUCompileMlir
   // CHECK-SAME: Y
   // CHECK: OpA
   // CHECK: OpB
   // CHECK: OpC
   "tf.OpA"() : () -> ()
   "tf_device.launch"() ({

  // CHECK: "tf.XlaHostCompute"(%arg0, %arg1)
  // CHECK-SAME: ancestors = []
  // CHECK-SAME: cost_estimate_ns = 1000000 : i64
  // CHECK-SAME: key = ""
  // CHECK-SAME: recv_key = "host_compute_channel_recv"
  // CHECK-SAME: send_key = "host_compute_channel_send"
  // CHECK-SAME: shape_inference_graph = @host_func
  // CHECK-SAME: shapes = [#tf_type.shape<*>, #tf_type.shape<3x?>]
  // CHECK-SAME: tpu_core = 0 : i64
  // CHECK: func @host_func

  // CHECK-SAME: ([[serialized]]
  // CHECK-NOT: device
  // CHECK-SAME: Tdense = [f32, f32]
  // CHECK-SAME: dense_shapes = [#corert.shape<>, #corert.shape<>]
  // CHECK-SAME: num_sparse = 2 : i64
  // CHECK-SAME: ragged_split_types = []
  // CHECK-SAME: ragged_value_types = []
  // CHECK-SAME: sparse_types = [!corert.string, i64]

  func.return %0 : tensor<i32>
}

// CHECK-LABEL: func private @merge_stateless_merged_if_0_0_then
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>)
// CHECK: [[r0:%.*]] = "tf.AddV2"([[x]], [[y]])
// CHECK: return [[r0]], [[r0]], [[r0]]

// CHECK-LABEL: func private @merge_stateless_merged_if_0_0_else
// CHECK-SAME: ([[x:%.*]]: tensor<i32>, [[y:%.*]]: tensor<i32>)

  // CHECK: "tf_saved_model.session_initializer"()
  // CHECK-SAME: initializers = [@init_func_restore_op]

  func.func @init_func_restore_op() -> () attributes {
      tf_saved_model.initializer_type = "restore_op",
      tf_saved_model.exported_names = ["tf_saved_model.session_initializer_restore_op"]} {
    return
  }
  // CHECK: @init_func_restore_op
  // CHECK-SAME: tf_saved_model.initializer_type = "restore_op"

}

// CHECK: func private @[[CLUSTER]]
// CHECK-SAME: (%[[CLUSTER_ARG_0:[a-z0-9]*]]: tensor<?xi32>) -> tensor<?xi32>
// CHECK: %[[B_OUTPUT:[0-9]*]] = "tf.B"(%[[CLUSTER_ARG_0]])
// CHECK: return %[[B_OUTPUT]]

// -----

// Tests that multiple `tf_device.cluster` that depend on each other are
// correctly handled.

// CHECK-LABEL: func @multiple_clusters
// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: tensor<?xi32>)

  func.return %1 : tensor<4x128x1xf32>
  // CHECK-NEXT: %[[CONST_WEIGHT:.*]] = arith.constant
  // CHECK-SAME: [1.000000e+00, 2.000000e+00]
  // CHECK-SAME: tensor<1x2xf32>
  // CHECK: %[[FC_RES:.*]] = "tfl.fully_connected"(%arg0, %[[CONST_WEIGHT]]
  // CHECK-SAME: <{fused_activation_function = "NONE", keep_num_dims = true, weights_format = "DEFAULT"}> : (tensor<4x128x2xf32>, tensor<1x2xf32>, none) -> tensor<4x128x1xf32>

// CHECK-LABEL: func @resource
func.func @resource() {
  // CHECK: tf.SomeOp
  // CHECK-SAME: dtype = !corert.resource
  "tf.SomeOp"() {device = "/device:CPU:0", dtype = !tf_type.resource} : () -> ()
  func.return
}

// CHECK-LABEL: func @variant
func.func @variant(%arg: tensor<!tf_type.variant>) {
  // CHECK: tf.ZerosLike
  // CHECK-SAME: T = !corert.variant

        then:
        dep.is(same)
        !dep.is(differentFrom)
        !dep.is(differentRequested)
        !dep.is(differentConstraint)
        dep.is(differentFailure) //the same dependency edge cannot have different failures
    }

    def selector(String group = 'a', String module = 'a', String version = '1') {

// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-unfuse-mhlo-batch-norm | FileCheck %s

// CHECK-LABEL: @unfuse_batch_norm
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
// CHECK-SAME: %[[MEAN:[^:[:space:]]+]]
// CHECK-SAME: %[[VARIANCE:[^:[:space:]]+]]
func.func @unfuse_batch_norm(
    %x: tensor<4x256xf32>, %scale: tensor<256xf32>, %offset: tensor<256xf32>,