%0 = "tf.BiasAdd"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", data_format = "NHWC"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32>
  func.return %0 : tensor<1x10x10x32xf32>
// CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%arg0, %arg1) <{data_format = "NHWC"}> {T = "tfdtype$DT_FLOAT"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32>

  }
}

Status ConvertScalarTypeToDataType(Type type, DataType* dtype) {
  if (type.isF16()) {
    *dtype = DT_HALF;
    return absl::OkStatus();
  } else if (type.isF32()) {
    *dtype = DT_FLOAT;
    return absl::OkStatus();
  } else if (type.isF64()) {
    *dtype = DT_DOUBLE;
    return absl::OkStatus();
  } else if (type.isBF16()) {
    *dtype = DT_BFLOAT16;

    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>

    %1 = "tf.LookupTableFindV2"(%0, %arg0, %cst) {device = ""} : (tensor<!tf_type.resource>, tensor<i64>, tensor<!tf_type.string>) -> tensor<*x!tf_type.string>
    func.return %1 : tensor<*x!tf_type.string>
    // CHECK-LABEL: @serving_default
    // CHECK: "tfl.call_once"() <{session_init_function = "init_all_tables"}> : () -> ()

  // CHECK-NEXT:  return [[FIND]] : tensor<*x!tf_type.string>
  func.return %1 : tensor<*x!tf_type.string>
}

// -----

// Test for case with unsupported string -> string mapping.
func.func @no_legalization_on_hashtable_string_to_string(%arg0: tensor<!tf_type.string>) {

func.func @variantAddN(%arg0: tensor<!tf_type.variant<tensor<*xi32>>>, %arg1: tensor<!tf_type.variant<tensor<*xi32>>>) -> tensor<!tf_type.variant<tensor<*xi32>>> {
  %1 = "tf.AddN"(%arg0, %arg1) : (tensor<!tf_type.variant<tensor<*xi32>>>, tensor<!tf_type.variant<tensor<*xi32>>>) -> tensor<!tf_type.variant<tensor<*xi32>>>

  %1 = "tf.MultiDeviceIteratorGetNextFromShard"(%0, %cst, %arg1) {device = ""} : (tensor<!tf_type.resource>, tensor<i32>, tensor<i64>) -> tensor<!tf_type.string>
  %2 = "tf.Identity"(%1) {device = ""} : (tensor<!tf_type.string>) -> tensor<!tf_type.string>
  return %2 : tensor<!tf_type.string>
}

func.func @quantized_types(%arg0: tensor<!tf_type.resource<tensor<1x3x!tf_type.quint8>>>,
                      %arg1: tensor<!tf_type.resource<tensor<1x3x!tf_type.quint16>>>,
                      %arg2: tensor<!tf_type.resource<tensor<1x3x!tf_type.qint8>>>,
                      %arg3: tensor<!tf_type.resource<tensor<1x3x!tf_type.qint16>>>,
                      %arg4: tensor<!tf_type.resource<tensor<1x3x!tf_type.qint32>>>) {
  // CHECK: tf.ReadVariableOp

    // 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>
    // CHECK-NEXT: return %[[UDQ]] : tensor<1xf32>

    %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = i64, shared_name = "hash_table_dba2ccaa-f1b1-46d6-b276-98008f69da71", use_node_name_sharing = false, value_dtype = !tf_type.string} : () -> tensor<!tf_type.resource>
    "tf.LookupTableImportV2"(%0, %cst, %cst_0) {device = ""} : (tensor<!tf_type.resource>, tensor<4xi64>, tensor<4x!tf_type.string>) -> ()