Type GetFixedElementType(StringRef element_type, Builder& builder) const {
    if (element_type == "i32_") return builder.getI32Type();
    if (element_type == "i64_") return builder.getI64Type();
    if (element_type == "f32_") return builder.getF32Type();
    if (element_type == "i1_") return builder.getI1Type();
    return {};
  }

tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>,...

      lhs_quantization_axis = -1 : i64,
      lhs_quantization_min_val = -2147483648 : i64,
      lhs_quantization_max_val = 2147483647 : i64,
      rhs_quantization_axis = -1 : i64,
      rhs_quantization_min_val = -2147483648 : i64,
      rhs_quantization_max_val = 2147483647 : i64,
      output_quantization_axis = -1 : i64,
      output_quantization_min_val = -2147483648 : i64,

    }, {
    ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
      %new_sum = "vhlo.add_v1"(%arg1, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
      %new_i = "vhlo.add_v1" (%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
      "vhlo.return_v1"(%new_i, %new_sum) : (tensor<i64>, tensor<i64>) ->()
  }) : (tensor<i64>, tensor<i64>) -> (tensor<i64>, tensor<i64>)
  return %0 : tensor<i64>
}

    device = "", lhs_quantization_axis = -1 : i64,
    lhs_quantization_max_val = 127 : i64,
    lhs_quantization_min_val = -128 : i64,
    output_quantization_axis = -1 : i64,
    output_quantization_max_val = 2147483647 : i64,
    output_quantization_min_val = -2147483648 : i64,
    rhs_quantization_axis = -1 : i64,
    rhs_quantization_max_val = 127 : i64,
    rhs_quantization_min_val = -128 : i64
  } : (

  %2 = "tf.Const"() {__op_key = 7: i32,device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
  %3 = "tf.RangeDataset"(%0, %2, %1) {__op_key = 8: i32, device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
  // CHECK: tf_mlrt.executeop{{.*}}op: \22MapDataset\22

    %4 = "tf.AddV2"(%2, %3) {device = ""} : (tensor<i64>, tensor<i64>) -> tensor<i64>
    %5 = "tf.LookupTableFindV2"(%0, %arg0, %cst) {device = ""} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
    %6 = "tf.AddV2"(%5, %4) {device = ""} : (tensor<*xi64>, tensor<i64>) -> tensor<*xi64>
    return %6 : tensor<*xi64>
  }
// Check that HashTable op is lifted.

// CHECK: %[[ADD:.+]] = stablehlo.add %[[CONV]], %[[BIAS_CONST]] : tensor<1x5x5x4xf32>
// CHECK: %[[REDUCE_WINDOW_MAX:.+]] = "stablehlo.reduce_window"(%[[ADD]], %[[INIT_VALUE_CONST:.+]])
// CHECK: <{window_dimensions = array<i64: 1, 2, 2, 1>, window_strides = array<i64: 1, 2, 2, 1>}>
// CHECK: stablehlo.maximum

"tf.BatchFunction"(%cst) <{allowed_batch_sizes = [1], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", enable_large_batch_splitting = true, f = @_batched, low_priority_allowed_batch_sizes = [], low_priority_batch_timeout_micros = 0 : i64, low_priority_max_batch_size = 0 : i64, low_priority_max_enqueued_batches = 0 : i64, max_batch_size = 1 : i64, max_enqueued_batches = 1 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch_function___inf...

      dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
      window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
      lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
      {batch_group_count = 1 : i64, feature_group_count = 1 : i64}
    : (tensor<?x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<?x2x2x1xi32>
  %1 = chlo.broadcast_add %0, %zp_offset : (
      tensor<?x2x2x1xi32>, tensor<?x2x2x1xi32>) -> tensor<?x2x2x1xi32>