"vhlo.return_v1"(%1112) : (tensor<i1>) -> ()
    }) : (tensor<448xf32>, tensor<448xi32>) -> (tensor<448xf32>, tensor<448xi32>)
  return %0 : tensor<448xf32>
}

//CHECK:func.func private @sort(%arg0: tensor<448xf32>, %arg1: tensor<448xi32>) -> tensor<448xf32> {
//CHECK-NEXT:  %0:2 = "vhlo.sort_v1"(%arg0, %arg1) <{dimension = #vhlo.integer_v1<0 : i64>, is_stable = #vhlo.bool_v1<true>}> ({

}

// -----

// Test invalid tf.Less
func.func @testLess(tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> {
^bb0(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>):
  // expected-error @+1 {{op result #0 must be tensor of bool values}}
  %0 = "tf.Less"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
  func.return %0 : tensor<4xi32>
}

// -----

// Test valid tf.ConcatV2

  %9 = mhlo.floor %8 : tensor<64xf32>
  %10 = mhlo.convert %9 : (tensor<64xf32>) -> tensor<64xi32>
  %11 = mhlo.compare  LT, %10, %1,  SIGNED : (tensor<64xi32>, tensor<64xi32>) -> tensor<64xi1>
  %12 = mhlo.add %10, %0 : tensor<64xi32>
  %13 = mhlo.select %11, %12, %10 : tensor<64xi1>, tensor<64xi32>
  %14 = mhlo.reshape %13 : (tensor<64xi32>) -> tensor<64x1xi32>

}

// -----

// CHECK-LABEL: func @testBroadcastGradientArgs
func.func @testBroadcastGradientArgs(%s0: tensor<4xi32>, %s1: tensor<4xi32>) -> (tensor<1xi32>, tensor<0xi32>) {
  // CHECK:     tf.BroadcastGradientArgs
  %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) : (tensor<4xi32>, tensor<4xi32>) -> (tensor<1xi32>, tensor<0xi32>)
  func.return %r0, %r1 : tensor<1xi32>, tensor<0xi32>
}

// -----

// CHECK-LABEL: @acos

  // Reduce over spatial dimensions (NCHW)
  %1 = "tf.Const"() {value = dense<[2, 3]> : tensor<2xi32>} : () -> tensor<2xi32>

  // Transpose input: NHWC -> NCHW
  %2 = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} : () -> tensor<4xi32>
  %3 = "tf.Transpose"(%arg0, %2) : (tensor<?x224x224x3xf32>, tensor<4xi32>) -> tensor<?x3x224x224xf32>

    >} : (
      tensor<?x?x6xi8>, tensor<6x8xi8>
    ) -> tensor<?x?x8xi32>
  %1 = chlo.broadcast_add %0, %zp_offset : (
      tensor<?x?x8xi32>, tensor<8xi32>) -> tensor<?x?x8xi32>
  %2 = chlo.broadcast_add %1, %bias : (
      tensor<?x?x8xi32>, tensor<8xi32>) -> tensor<?x?x8xi32>
  return %2 : tensor<?x?x8xi32>
}

// -----

// CHECK-LABEL: func @dot_general_add_add_static

     %6 = "tf.Identity"(%4) : (tensor<2048xi32>) -> tensor<2048xi32>
    "tf.TPUExecute"(%6, %arg0) {device = "foobar"} : (tensor<2048xi32>, tensor<!tf_type.string>) -> ()
    return
}

// -----

// CHECK-LABEL test_disconnected
func.func @test_disconnected(%arg0: tensor<!tf_type.string>) {
     %0 = builtin.unrealized_conversion_cast to tensor<2048xi32>
     %1 = builtin.unrealized_conversion_cast to tensor<2048xi32>

	  tf_device.return %3#0, %3#1 : tensor<2048xi32>, tensor<2048xi32>
	}) {device = "TPU_REPLICATED_HOST_0"} : () -> (tensor<2048xi32>, tensor<2048xi32>)
    return %0#0, %0#1: tensor<2048xi32>, tensor<2048xi32>
  }

  // CHECK-LABEL: func @valid_copy_op_in_non_replicated_host

  // CHECK: "tf_device.launch"

// CHECK-NEXT:  %[[argmax:.*]] = "tfl.arg_max"(%[[softmax]], %[[cst]]) : (tensor<128x16xf32>, tensor<i32>) -> tensor<128xi32>
  %3 = "tfl.arg_max"(%2, %cst) : (tensor<128x16xf32>, tensor<i32>) -> tensor<128xi32>
// CHECK-NEXT:  return %[[softmax]], %[[argmax]] : tensor<128x16xf32>, tensor<128xi32>
  func.return %2, %3 : tensor<128x16xf32>, tensor<128xi32>
}

// CHECK-LABEL: PruneUnusedLstm

            // CHECK-NOT: tf.TPUAnnotateTensorsWithDynamicShape
            %4:2 = "tf.TPUAnnotateTensorsWithDynamicShape"(%3#0, %3#1) : (tensor<2048xi32>, tensor<2048xi32>) -> (tensor<2048xi32>, tensor<2048xi32>)
            tf_device.return %4#0, %4#1 : tensor<2048xi32>, tensor<2048xi32>