%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>

  func.return %result : tensor<1x1x512xf32>

module {
func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
  %1 = arith.constant dense<1> : tensor<1xi32>
  %2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
  %3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
    : (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>
  func.return %3 : tensor<2x1x2xf32>
}
}

//CHECK: module attributes

module {
func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
  %1 = arith.constant dense<1> : tensor<1xi32>
  %2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
  %3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
    : (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>
  func.return %3 : tensor<2x1x2xf32>
}
}

//CHECK: module {

func.func @mean_4d_no_keepdim(%arg0: tensor<1x48x48x512xf32>) -> tensor<1x512xf32> {
  %cst = arith.constant dense<[1, 2]> : tensor<2xi32>
  %0 = "tfl.mean"(%arg0, %cst) {keep_dims = false} : (tensor<1x48x48x512xf32>, tensor<2xi32>) -> tensor<1x512xf32>
  func.return %0 : tensor<1x512xf32>
}

// CHECK:       func @mean_4d_no_keepdim([[VAL_0:%.*]]: tensor<1x48x48x512xf32>) -> tensor<1x512xf32> {

  %cst_add = arith.constant dense<2.0> : tensor<2x512xf32>  // Not 1D
  %cst_weights = arith.constant dense<3.0> : tensor<1024x512xf32>
  %cst_bias = arith.constant dense<5.0> : tensor<1024xf32>

  %0 = "tfl.add"(%arg, %cst_add) {fused_activation_function = "NONE"} : (tensor<2x512xf32>, tensor<2x512xf32>) -> tensor<2x512xf32>

  %cst_input = arith.constant dense<1.0> : tensor<2x512xf32>
  %cst_weights = arith.constant dense<2.0> : tensor<1024x512xf32>
  %cst_bias = arith.constant dense<4.0> : tensor<1024xf32>

  %0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x512xf32>, tensor<1024x512xf32>, tensor<1024xf32>) -> tensor<2x1024xf32>

}

func.func @squeezeSingleAxis(%arg0: tensor<2x1x2xf32>) -> tensor<2x2xf32> {
  %0 = "tf.Squeeze"(%arg0) {squeeze_dims = [1]} : (tensor<2x1x2xf32>) -> tensor<2x2xf32>
  func.return %0 : tensor<2x2xf32>

// CHECK-LABEL:squeezeSingleAxis
// CHECK:  "tfl.squeeze"(%arg0) <{squeeze_dims = [1]}> : (tensor<2x1x2xf32>) -> tensor<2x2xf32>
}

func.func @quantized_dot_general_not_converted(%arg0: tensor<1x1x512xf32>, %arg1: tensor<512x512x!quant.uniform<i8:f32, 0.00285>>) -> tensor<1x1x512xf32> {
  %0 = "mhlo.dot_general"(%arg0, %arg1) {
    dot_dimension_numbers = #mhlo.dot<
      lhs_contracting_dimensions = [2],
      rhs_contracting_dimensions = [0]
    >} : (tensor<1x1x512xf32>, tensor<512x512x!quant.uniform<i8:f32, 0.00285>>) -> tensor<1x1x512xf32>
  func.return %0 : tensor<1x1x512xf32>
}

    %cst_5 = arith.constant dense<512> : tensor<i32>
    %19 = "tf.Pack"(%18, %cst_5) {N = 2 : i64, T = i32, axis = 0 : i64, device = ""} : (tensor<i32>, tensor<i32>) -> tensor<2xi32>
    // CHECK: -> tensor<1x512xf32>
    %40 = "tf.Reshape"(%39, %19) {T = f32, Tshape = i32, device = ""} : (tensor<1x4x4x32xf32>, tensor<2xi32>) -> tensor<?x?xf32>
   func.return
  }

  // CHECK-LABEL: func @shape_partial_eval