%0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i64>} : () -> tensor<i64>
    %1 = "tf.Mul"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i64>, tensor<i64>) -> tensor<i64>
    %2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<i64>) -> tensor<i64>
    func.return %2 : tensor<i64>
  }

    // CHECK-NEXT: %[[UDQ:.*]] = "tf.UniformDequantize"(%[[UQ]], %[[CST_0]], %[[CST_1]]) <{quantization_axis = -1 : i64, quantization_max_val = 127 : i64, quantization_min_val = -128 : i64}> : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>

  "tf.Assert"(%cond, %error_message) {device = "/job:localhost/replica:0/task:0/device:CPU:0", summarize = 3 : i64} : (tensor<i1>, tensor<!tf_type.string>) -> ()
  %value1 = "tf.LookupTableFindV2"(%handle, %key1, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<!tf_type.string>, tensor<i64>) -> tensor<i64>

    %1 = "tf.Const"() {device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
    %2 = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>

    %0 = stablehlo.composite "stablehlo.add_n" %arg0 { composite_attributes = { test_int = 2 : i64, test_bool = 0 : i1, test_string = "test"}, decomposition = @add_n.impl } : (tensor<i64>) -> tensor<i64>
    return %0 : tensor<i64>
  }
  func.func private @add_n.impl(%arg0: tensor<i64>) -> tensor<i64> {
    %0 = stablehlo.constant dense<2> : tensor<i64>

      %arg1: tensor<i64>
    ) {
    // CHECK: tf.ReduceDataset
    %1 = "tf.ReduceDataset"(%arg0, %arg1) {
      Targuments = [],
      Tstate = [i64], device = "",
      f = @__reduce_func_0, f._tf_data_function = true,
      output_shapes = [#tf_type.shape<>],
      output_types = [i64], use_inter_op_parallelism = true } : (tensor<!tf_type.variant>, tensor<i64>) -> (tensor<i64>)
    func.return
}

          tf_device.return
        }) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> ()
        "tf_device.launch"() ({
          %1 = "tf.opA"() : () -> tensor<!tf_type.string>
          "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"(%1){_xla_replica_id = 1 : i64, device_ordinal = -1 : i64} : (tensor<!tf_type.string>) -> ()

  %1 = "tf.LookupTableSizeV2"(%0) {device = ""} : (tensor<!tf_type.resource>) -> tensor<i64>
  %2 = "tf.LookupTableFindV2"(%0, %arg, %default) {device = "/CPU:0"} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
  func.return %1, %2 : tensor<i64>, tensor<*xi64>
}

}

  // CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""
  // CHECK-SAME: enable_large_batch_splitting = false
  // CHECK-SAME: max_batch_size = 6 : i64
  // CHECK-SAME: max_enqueued_batches = 10 : i64
  // CHECK-SAME: num_batch_threads = 2 : i64
  // CHECK-SAME: shared_name = "batch/"

  // CHECK-NEXT:  return [[SIZE]] : tensor<i64>
  func.return %1 : tensor<i64>
}

// -----

// Test for case with import and find ops.
func.func @hashtable_import_then_find(%arg0: tensor<5x!tf_type.string>) -> tensor<*xi64> {
  %cst = arith.constant dense<["emerson", "lake", "palmer"]> : tensor<3x!tf_type.string>
  %cst_0 = arith.constant dense<-1> : tensor<i64>