func.return %0 : tensor<4x128xf32>
  }
  func.func @_func(%arg0: tensor<4x128xf32>) -> tensor<4x128xf32> {
    func.return %arg0 : tensor<4x128xf32>
  }

// CHECK:           %[[VAL_5:.*]] = "tf.Select"(%[[VAL_3]], %[[VAL_0]], %[[VAL_4]]) : (tensor<?x?xi1>, tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32>
// CHECK:           return %[[VAL_5]] : tensor<4x8xf32>
// CHECK:         }
func.func @relu_grad(%arg0: tensor<4x8xf32>, %arg1: tensor<?x?xf32>) -> tensor<4x8xf32> {

func.func @testConv2dShapeInferenceDynamic(%arg0: tensor<1x?x?x128xf32>, %arg1: tensor<128x3x3x128xf32>, %arg2: tensor<128xf32>) -> tensor<1x?x?x128xf32> {
  // CHECK: "tfl.conv_2d"(%arg0, %arg1, %arg2) <{dilation_h_factor = 2 : i32, dilation_w_factor = 2 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x?x?x128xf32>, tensor<128x3x3x128xf32>, tensor<128xf32>) -> tensor<1x?x?x128xf32>

  func.return %4 : tensor<8x128xf32>

// CHECK-LABEL: SoftMaxWithNormalization
// CHECK: %[[RESULT:.*]] = "tfl.softmax"(%arg0) <{beta = 1.000000e+00 : f32}> : (tensor<8x128xf32>) -> tensor<8x128xf32>
// CHECK: return %[[RESULT]] : tensor<8x128xf32>
}

func.func @SoftMaxWithoutNormalization(%arg0: tensor<8x128xf32>) -> tensor<8x128xf32> {

  func.return %1 : tensor<4x128x1xf32>

  // CHECK: %[[TRANSPOSED_X:.*]] = "tfl.transpose"
  // CHECK-SAME: (tensor<4x2x128xf32>, tensor<3xi32>) -> tensor<4x128x2xf32>
  // CHECK-NEXT: %[[FC_RES:.*]] = "tfl.fully_connected"(%[[TRANSPOSED_X]]
  // CHECK-SAME: <{fused_activation_function = "NONE", keep_num_dims = true, weights_format = "DEFAULT"}> : (tensor<4x128x2xf32>, tensor<1x2xf32>, none) -> tensor<4x128x1xf32>

    %12 = "tfl.logistic"(%11) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %13 = tfl.mul %arg4, %12 {fused_activation_function = "NONE"} : tensor<4x2xf32>
    %14 = "tfl.relu"(%10#1) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %15 = "tfl.logistic"(%10#0) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %16 = tfl.mul %15, %14 {fused_activation_function = "NONE"} : tensor<4x2xf32>

func.func @testAddV2IdentityBroadcastTensor(%arg0: tensor<4x1xf32>, %arg1: tensor<4x2xf32>) -> (tensor<4x2xf32>, tensor<4x2xf32>, tensor<4x2xf32>, tensor<4x2xf32>) {
  %0 = "tf.Const"() {value = dense<0.0> : tensor<1x2xf32>} : () -> tensor<1x2xf32>

  // Operand and identity shapes are broadcastable. However, we cannot fold
  // because the operand does not match the result shape.
  %1 = "tf.AddV2"(%arg0, %0) : (tensor<4x1xf32>, tensor<1x2xf32>) -> tensor<4x2xf32>

  %w = arith.constant dense<127.0> : tensor<4x12xf32>
  %b = arith.constant dense<0.0> : tensor<4xf32>
  %fc = "tfl.fully_connected"(%arg0, %w, %b) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x224x224x3xf32>, tensor<4x12xf32>, tensor<4xf32>) -> tensor<1x112x112x4xf32>
  func.return %fc : tensor<1x112x112x4xf32>

    // CHECK-NEXT: return %[[DIV]] : tensor<4x10xf32>
    %0 = "tf.Softsign"(%arg0) : (tensor<4x10xf32>) -> tensor<4x10xf32>
    func.return %0 : tensor<4x10xf32>
}

// -----

// CHECK-LABEL: func @softsign_grad
func.func @softsign_grad(%arg0: tensor<4x10xf32>, %arg1: tensor<4x10xf32>) -> tensor<4x10xf32> {

    // CHECK-NEXT: %[[ONE:.*]] = mhlo.constant dense<1.000000e+00> : tensor<f32>

// CHECK:           %[[VAL_14:.*]] = "tfl.reshape"(%[[VAL_13]], %[[VAL_8]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x384x128xf32>, tensor<2xi32>) -> tensor<384x128xf32>
// CHECK:           %[[VAL_15:.*]] = "tfl.reshape"(%[[VAL_14]], %[[VAL_5]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<384x128xf32>, tensor<3xi32>) -> tensor<1x384x128xf32>