// Tests transitive Identity ops reachable from the cluster are pruned away.

// CHECK-LABEL: func @testTransitiveIdentity
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>)
func.func @testTransitiveIdentity(%arg0: tensor<i32>) {
  // CHECK:      "tf_device.cluster"
  // CHECK:      "tf.PartitionedCall"([[ARG0]])
  // CHECK-SAME: f = @callee0
  %0 = "tf_device.cluster"() ({

  }
}

// -----

module attributes {tf_saved_model.semantics} {
  // CHECK-LABEL: composite_and_select
  // CHECK-SAME: %[[ARG0:.*]]: tensor<1x2xf32>
  // CHECK-SAME: %[[ARG1:.*]]: tensor<2x3xf32>
  // CHECK-SAME: %[[ARG2:.*]]: tensor<1x3xi1>
  // CHECK-SAME: %[[ARG3:.*]]: tensor<1x3xf32>

// CHECK-LABEL: func private @composite_depthwise_conv2d_fn(
// CHECK-SAME:                                             %arg0: tensor<1x3x4x3xf32>,
// CHECK-SAME:                                             %arg1: tensor<2x3x3x2xf32>)

// CHECK-LABEL: func private @composite_depthwise_conv2d_fn_0(
// CHECK-SAME:                                             %arg0: tensor<1x3x4x3xf32>,
// CHECK-SAME:                                             %arg1: tensor<2x3x1x6xf32>)

// CHECK-SAME: %[[ARG_0:[a-z0-9]*]]: tensor<?x2xi32>
// CHECK-SAME: %[[ARG_1:[a-z0-9]*]]: tensor<?x2xi32>
// CHECK-SAME: %[[ARG_2:[a-z0-9]*]]: tensor<?x2xi32>
// CHECK-SAME: %[[ARG_3:[a-z0-9]*]]: tensor<?xi32>
// CHECK-SAME: %[[ARG_4:[a-z0-9]*]]: tensor<?xi32>
// CHECK-SAME: %[[ARG_5:[a-z0-9]*]]: tensor<?xi32>
// CHECK-SAME: %[[ARG_6:[a-z0-9]*]]: tensor<!tf_type.string>
// CHECK-SAME: %[[ARG_7:[a-z0-9]*]]: tensor<!tf_type.string>

// CHECK-LABEL: func @main
// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 0 : i32}

// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 1 : i32}

// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 2 : i32}

// CHECK: tf.AddV2
// CHECK-SAME: {__op_key = 3 : i32}

// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 4 : i32}

// CHECK: tf.Sub
// CHECK-SAME: {__op_key = 5 : i32}

// CHECK: tf.Sub

# it is being invoked.
# CHECK: module
# CHECK-SAME: tf.versions
# CHECK-SAME: bad_consumers
# CHECK-SAME: min_consumer
# CHECK-SAME: producer

# CHECK: "tf_saved_model.global_tensor"() <{is_mutable, sym_name = "[[VAR:[a-zA-Z_0-9]+]]", type = tensor<1x3xf32>, value = {{.*}} : tensor<1x3xf32>}> : () -> ()

# CHECK:      func {{@[a-zA-Z_0-9]+}}(
# CHECK-SAME:   [[ARG0:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]},

            }
        }
    }

    @Test
    fun `same script different target type`() {

        // given: same init, settings & build files all applying same script
        val same = withFile("same.gradle.kts", """println("Same script on ${'$'}this")""")
        val sameApply = """apply(from = "same.gradle.kts")"""
        val initScriptFile = withFile("same.init.gradle.kts", sameApply)

// CHECK-LABEL: @batchNormInference_2D_inner_features
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
// CHECK-SAME: %[[MEAN:[^:[:space:]]+]]
// CHECK-SAME: %[[VARIANCE:[^:[:space:]]+]]
func.func @batchNormInference_2D_inner_features(

// CHECK:       ENTRY %main.{{[0-9]+}}
// CHECK-SAME:  (arg_tuple.{{[0-9]+}}: (f32[128,10], f32[10,1024], f32[128,1024])) -> (f32[128,10], f32[10,1024], f32[128,1024]) {
// CHECK:         ROOT %tuple.{{[0-9]+}}
// CHECK-SAME:    sharding={
// CHECK-SAME:    {devices=[1,2]0,1}
// CHECK-SAME:    {maximal device=0}
// CHECK-SAME:    {replicated}

  #
  # CHECK:      func {{@[a-zA-Z_0-9]+}}(
  # CHECK-SAME:   %arg0: tensor<f32> {tf._user_specified_name = "x", tf_saved_model.index_path = [0]},
  # CHECK-SAME:   %arg1: tensor<!tf_type.resource<{{.*}}>>
  # CHECK-SAME: ) -> (
  # CHECK-SAME:   tensor<f32> {tf_saved_model.index_path = [0]},
  # CHECK-SAME:   tensor<f32> {tf_saved_model.index_path = [1]})
  # CHECK-SAME: attributes{{.*}}tf_saved_model.exported_names = ["callee"]