// CHECK-LABEL: func @transpose_resnet_layer
func.func @transpose_resnet_layer(%arg0: tensor<?x224x224x3xf32>, // input
                             %arg1: tensor<64xf32>,          // batch_norm args
                             %arg2: tensor<256xf32>,          // batch_norm args
                             %arg3: tensor<7x7x3x64xf32>,    // conv filter #0
                             %arg4: tensor<1x1x64x256xf32>   // conv filter #1

// -----

module {
func.func @inference_can_fuse(%arg0: tensor<?x8x8xf32>, %arg1: tensor<8x10xf32>, %arg2: tensor<8x10xf32>, %arg3: tensor<8x40xf32>, %arg4: tensor<10x40xf32>, %arg5: tensor<40xf32>) {
  %0 = "tf.Const"() {_output_shapes = ["tfshape$"], device = "", dtype = f32, value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>

// CHECK-LABEL: func @do_not_hoist_writes_to_explicitly_captured_resources
// CHECK-SAME:  ([[ARG0:%.*]]: tensor<*x!tf_type.resource<tensor<f32>>>, [[ARG1:%.*]]: tensor<*x!tf_type.resource<tensor<f32>>>) {
func.func @do_not_hoist_writes_to_explicitly_captured_resources(%arg0: !tf_res_f32, %arg1: !tf_res_f32) {
  // CHECK: [[REPLICATE:%.*]]:2 = tf_device.replicate({{\[}}[[ARG0]], [[ARG1]]] as [[RI:%.*]]: tensor<*x!tf_type.resource<tensor<f32>>>) {n = 2 : i32}

func.func @replicate_with_packed_input(%arg0: tensor<i1>, %arg1: tensor<i1>) {
  %0:4 = tf_executor.graph {
    %1:5 = tf_executor.island {
      %2:4 = tf_device.replicate(%arg0 as %arg2: tensor<i1>, %arg1 as %arg3: tensor<i1>)
          {n = 2 : i32, _packed_input_indices = [0, 1]} {
        %3 = "tf.opA"(%arg2) : (tensor<i1>) -> tensor<f32>
        %4 = "tf.opB"(%arg3) : (tensor<i1>) -> tensor<i32>

# CHECK-SAME: attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "arg0,arg1,arg2", outputs = "ret0,ret1,ret2"}}...

  }

  func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
    return %0 : tensor<1x3xf32>
  }
}

// CHECK-LABEL: quantize_dot_general_fn
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>

}

func.func @branch0(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
  %0 = "tf.Add" (%arg0, %arg1) {device = "/device:CPU:0"}  : (tensor<f32>, tensor<f32>) -> tensor<f32>
  func.return %0 : tensor<f32>
}

func.func @branch1(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
  %0 = "tf.Add" (%arg0, %arg1) {device = "/device:CPU:0"}  : (tensor<f32>, tensor<f32>) -> tensor<f32>

func.func @xla_dot_v2(%arg0: tensor<?x2x3xf32>, %arg1: tensor<3x4x5xf32>) -> (tensor<?x2x4x5xf32>) {
  %0 = "tf.XlaDotV2"(%arg0, %arg1) {device = "", dimension_numbers = "\0A\01\02\12\01\00", precision_config = ""} : (tensor<?x2x3xf32>, tensor<3x4x5xf32>) -> tensor<?x2x4x5xf32>
  func.return %0 : tensor<?x2x4x5xf32>
}

// CHECK: func @xla_dot_v2

  func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>, %arg2: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
    %1 = stablehlo.add %0, %arg2 : tensor<1x3xf32>
    return %0 : tensor<1x3xf32>
  }

// CHECK: func.func private @serving_default
// CHECK-SAME: (%arg0: tensor<?x!tf_type.string>, %arg1: tensor<!tf_type.resource>) -> tensor<*xi64>
// CHECK-SAME: tf.entry_function = {control_outputs = "", inputs = "input_vocabs:0,hash_table_1:0", outputs = "FakeQuantWithMinMaxArgs_2:0"}
// CHECK: "tf.PartitionedCall"(%arg0, %arg1)
  func.func private @serving_default1(%arg0: tensor<?x!tf_type.string> ) -> (tensor<*xi64>) {