%2 = "mhlo.slice"(%arg0) <{limit_indices = dense<[3, 1, 512]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}> : (tensor<3x1x512xf32>) -> tensor<1x1x512xf32>
  %r = "mhlo.concatenate"(%0, %1, %2) <{dimension = 0 : i64}> : (tensor<1x1x512xf32>, tensor<1x1x512xf32>, tensor<1x1x512xf32>) -> tensor<3x1x512xf32>

  %0 = "tfl.fully_connected"(%arg0, %arg1, %arg2) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<384x384xf32>, tensor<512x384xf32>, tensor<512xf32>) -> tensor<384x512xf32>
  func.return %0: tensor<384x512xf32>
}

  %w = arith.constant dense<[[[[-1.0, 1.0]]], [[[1.0, 2.0]]], [[[-2.0, 1.0]]]]> : tensor<3x1x1x2xf32>
  %b = arith.constant dense<[1.0e-2, 2.1473647e1, -2.1473647e2]> : tensor<3xf32>
  %transpose_conv = "tfl.transpose_conv"(%arg1, %w, %0, %b) {
    padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32, fused_activation_function = "NONE"
  } : (tensor<4xi32>, tensor<3x1x1x2xf32>, tensor<1x5x5x2xf32>, tensor<3xf32>) -> tensor<1x5x5x3xf32>

// CHECK:           %[[VAL_10:.*]] = "tfl.dequantize"(%[[VAL_1]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.000000e-01>>) -> tensor<128x512xf32>
// CHECK:           %[[VAL_11:.*]] = "tfl.reshape"(%[[VAL_9]], %[[VAL_6]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<384x512xf32>, tensor<4xi32>) -> tensor<1x1x384x512xf32>

  %0 = "quantfork.stats"(%arg0) {
    layerStats = dense<[0.0, 1.0]> : tensor<2xf32>
  } : (tensor<?x5x5x2xf32>) -> tensor<?x5x5x2xf32>
  %1 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<3x5x5x2xf32>} : () -> tensor<3x5x5x2xf32>
  %2 = "tfl.pseudo_const"() {value = dense<0.000000e+00> : tensor<3xf32>} : () -> tensor<3xf32>
  %3 = "tfl.conv_2d"(%0, %1, %2) {
    dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32,

//     {0}, then it returns: tensor<1x3x5xi32>.
//   * If `xla_gather_op_output_type` == tensor<3x5xf32> and `collapsed_dims` ==
//     {1, 3}, then it returns: tensor<3x1x5x1xf32>.
Type GetSliceOpOutputType(Type xla_gather_op_output_type,
                          const absl::flat_hash_set<int64_t>& collapsed_dims) {
  if (auto ranked_output_type =

// CHECK-NEXT:    %[[RESHAPED_1:.*]] = mhlo.reshape %[[TRANSPOSED_1]]
// CHECK-NEXT:    %[[BMM_0:.*]] = "tfl.batch_matmul"(%[[RESHAPED_0]], %[[RESHAPED_1]]) <{adj_x = false, adj_y = false, asymmetric_quantize_inputs = false}> : (tensor<3x5x12xf32>, tensor<3x12x4xf32>) -> tensor<3x5x4xf32>
// CHECK-NEXT:    %[[RESHAPED_BMM:.*]] = mhlo.reshape %[[BMM_0]]
// CHECK-NEXT:    return %[[RESHAPED_BMM]] : tensor<3x5x1x4xf32>
}