// CHECK: %0 = tfl.add(%arg0, %cst_0) <{fused_activation_function = "NONE"}> : (tensor<2x3x4x5xf32>, tensor<2x3x4x1xf32>) -> tensor<2x3x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>

// -----

// CHECK-LABEL: pushTposeBcastScalarCstInput
func.func @pushTposeBcastScalarCstInput(%arg0: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {

  return %0 : tensor<1x1x1x5xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_6(%arg0: tensor<1x1x1x8xf32>) -> tensor<1x1x1x5xf32>

// CHECK-LABEL: avg_pool2d_7
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>

// -----

// CHECK-LABEL: transpose_simple_4d
func.func @transpose_simple_4d() -> tensor<5x2x3x4xf32> {
  %0 = stablehlo.constant dense<1.000000e+0> : tensor<2x3x4x5xf32>
  %1 = stablehlo.transpose %0, dims = [3, 0, 1, 2] : (tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32>
  return %1 : tensor<5x2x3x4xf32>
}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<5x2x3x4xf32>
// CHECK-NOT: transpose

func.func @ConvertSliceToIdentityI32(%arg0: tensor<2x3x4x5xf32>) -> tensor<2x3x4x5xf32> {
  %begin = arith.constant dense<0> : tensor<4xi32>
  %shape = arith.constant dense<[2,3,4,5]> : tensor<4xi32>
  %0 = "tfl.slice"(%arg0, %begin, %shape) : (tensor<2x3x4x5xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<2x3x4x5xf32>
  func.return %0 : tensor<2x3x4x5xf32>
  // CHECK: return %arg0
}

  return %2 : tensor<?x2x2x1xi8>
}

// -----

// CHECK-LABEL: func @convolution_add_add_f32
func.func @convolution_add_add_f32(
    %lhs: tensor<?x3x2x1xf32>, %rhs: tensor<2x1x1x1xf32>,
    %zp_offset: tensor<?x2x2x1xf32>, %bias: tensor<1xf32>
  ) -> tensor<?x2x2x1xf32> {
  // CHECK-DAG: %[[conv:.*]] = mhlo.convolution
  // CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[conv:.*]], %[[zp_offset:.*]]

  %cst_1 = arith.constant dense<6.000000e+00> : tensor<2x1x3x3xf32>
  %0 = "tfl.quantize"(%cst_1) {qtype = tensor<2x1x3x3x!quant.uniform<i8<-127:127>:f32:0, {6.587140e-03,1.888450e-02}>>} : (tensor<2x1x3x3xf32>) -> tensor<2x1x3x3x!quant.uniform<i8<-127:127>:f32:0, {6.587140e-03,1.888450e-02}>>
  %1 = "tfl.dequantize"(%0) : (tensor<2x1x3x3x!quant.uniform<i8<-127:127>:f32:0, {6.587140e-03,1.888450e-02}>>) -> tensor<2x1x3x3xf32>

  auto build_result = BuildHloFromGraph(*graph, /*use_output_shapes=*/true);
  ASSERT_FALSE(build_result.ok());
  EXPECT_THAT(build_result.message(),
              HasSubstr("op operand type 'tensor<2x3x4x5xi32>' and result type "
                        "'tensor<1xi32>' are cast incompatible"));
}

}  // namespace

// Tests that a float `stablehlo.constant` is not converted into `tfl.qconst`.

// CHECK-LABEL: func @float_constant
func.func @float_constant() -> tensor<1x2x4x5xf32> {
  %0 = stablehlo.constant() {value = dense<1.0> : tensor<1x2x4x5xf32>} : () -> tensor<1x2x4x5xf32>
  return %0 : tensor<1x2x4x5xf32>
}

// CHECK: stablehlo.constant
// CHECK-NOT: tfl.pseudo_qconst
// CHECK-NOT: tfl.pseudo_const
// CHECK-NOT: arith.constant

}

func.func @LRN(%arg0: tensor<2x3x4x5xf32>) -> tensor<2x3x4x5xf32> {
  %0 = "tf.LRN"(%arg0) {depth_radius = 5 :i64, bias = 1.0 :f32, alpha = 1.0 : f32, beta = 0.5 :f32} : (tensor<2x3x4x5xf32>) -> (tensor<2x3x4x5xf32>)
  func.return %0: tensor<2x3x4x5xf32>

  // CHECK-LABEL: LRN

  %0 = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} : () -> tensor<4xi32>
  %1 = "tf.Const"() {value = dense<[2, 0, 3, 1]> : tensor<4xi32>} : () -> tensor<4xi32>
  %2 = "tf.Transpose"(%arg0, %0) : (tensor<1x4x4x8xf32>, tensor<4xi32>) -> tensor<1x8x4x4xf32>
  %3 = "tf.Transpose"(%2, %1) : (tensor<1x8x4x4xf32>, tensor<4xi32>) -> tensor<4x1x4x8xf32>

  func.return %3 : tensor<4x1x4x8xf32>