// CHECK-NEXT: }) : (tensor<2x3x4xui8>) -> tensor<2x3x4xui8>
  // CHECK-NEXT: return %[[SORT]]
  %0 = "tf.XlaVariadicSort"(%arg0, %cst) {_XlaHasReferenceVars = false, comparator = @compare_lt, device = "/job:localhost/replica:0/task:0/device:XLA_GPU:0", is_stable = false} : (tensor<2x3x4xui8>, tensor<i32>) -> tensor<2x3x4xui8>
  func.return %0 : tensor<2x3x4xui8>
}

  // CHECK: %[[BROADCAST:.*]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<4xui8>) -> tensor<2x4xui8>
  // CHECK:  mhlo.shift_right_logical %[[BROADCAST]], %arg1 : tensor<2x4xui8>
  %0 = "tf.RightShift"(%arg0, %arg1) : (tensor<4xui8>, tensor<2x4xui8>) -> tensor<2x4xui8>
  func.return %0 : tensor<2x4xui8>
}

    %5 = "tf.Cast"(%4) {Truncate = false, device = ""} : (tensor<2x3x4xf32>) -> tensor<2x3x4xi8>
    %6 = "tf.Identity"(%5) {device = ""} : (tensor<2x3x4xi8>) -> tensor<2x3x4xi8>
    %7 = "tf.Cast"(%6) {Truncate = false, device = ""} : (tensor<2x3x4xi8>) -> tensor<2x3x4xi32>
    %8 = "tf.Sub"(%7, %cst_4) {device = ""} : (tensor<2x3x4xi32>, tensor<i32>) -> tensor<2x3x4xi32>

  func.return %0 : tensor<*xf32>
}

// -----

// CHECK-LABEL: func @testReverseV2
func.func @testReverseV2(%arg0: tensor<2x4x3xui8>, %arg1: tensor<1xi32>) -> tensor<2x4x3xui8> {
  %0 = "tf.ReverseV2"(%arg0, %arg1) : (tensor<2x4x3xui8>, tensor<1xi32>) -> tensor<2x4x3xui8>
  func.return %0 :  tensor<2x4x3xui8>
}

// -----

func.func @testIdentityWrongType(%arg0: tensor<4x2x!tf_type.string>) -> tensor<4x2x!tf_type.stringref> {