// -----

// Tests that int constants are not folded.

// CHECK-LABEL: transpose_int
func.func @transpose_int() -> tensor<3x2xi32> {
  %0 = stablehlo.constant dense<0> : tensor<2x3xi32>
  %1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2x3xi32>) -> tensor<3x2xi32>
  return %1 : tensor<3x2xi32>
}
// CHECK: transpose

// -----

// Tests that transposing an argument cannot be folded.

// CHECK-LABEL: transpose_arg

  func.func private @quantize_gather_fn(%arg: tensor<3x4x2xf32>) -> tensor<2x3x2x2xf32> attributes {tf._original_func_name = "main_0"} {
    %cst = "tf.Const"() {value = dense<1> : tensor<2x3x2xi32>} : () -> tensor<2x3x2xi32>
    %0 = "quantfork.stats"(%arg) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<3x4x2xf32>) -> tensor<3x4x2xf32>

    %7 = stablehlo.broadcast_in_dim %6, dims = [2, 1] : (tensor<1x3xf32>) -> tensor<2x3x2xf32>
    %8 = "quantfork.qcast"(%7) {volatile} : (tensor<2x3x2xf32>) -> tensor<2x3x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
    %9 = "quantfork.dcast"(%8) : (tensor<2x3x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2xf32>
    return %9 : tensor<2x3x2xf32>
  }

  // CHECK: quantized_dot_general_fn_1

  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

}

// -----

// CHECK-LABEL: func @select_batch_static_r1
func.func @select_batch_static_r1(%arg0: tensor<i1>, %arg1: tensor<2x6x8xi32>, %arg2: tensor<2x6x8xi32>) -> tensor<2x6x8xi32> {
  // CHECK: mhlo.select %arg0, %arg1, %arg2
  %0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<2x6x8xi32>, tensor<2x6x8xi32>) -> tensor<2x6x8xi32>
  func.return %0: tensor<2x6x8xi32>
}

// -----

// Examples:
//   * If `xla_gather_op_output_type` == tensor<*xf32>, then it returns:
//     tensor<*xf32>.
//   * If `xla_gather_op_output_type` == tensor<3x5xi32> and `collapsed_dims` ==
//     {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,

func.func @ConvertIdentityGatherNdOp3D(%arg0: tensor<4x3x4xf32>) -> tensor<4x3x4xf32> {
  %cst = arith.constant dense<[[0], [1], [2], [3]]> : tensor<4x1xi32>
  %0 = "tfl.gather_nd"(%arg0, %cst) : (tensor<4x3x4xf32>, tensor<4x1xi32>) -> tensor<4x3x4xf32>
  func.return %0 : tensor<4x3x4xf32>

// CHECK-LABEL: ConvertIdentityGatherNdOp3D
// CHECK-SAME: (%[[ARG:.*]]: tensor<4x3x4xf32>) -> tensor<4x3x4xf32>

    %0 = stablehlo.constant dense<3.000000e+00> : tensor<1x1xf32>  // Input inverse scale.
    %1 = stablehlo.constant dense<1> : tensor<1x1xi8>  // Input zero point.
    %2 = stablehlo.constant dense<5> : tensor<2x3xi32>  // Quantized filter - the pattern expects i8 but i32 is given.
    %3 = stablehlo.constant dense<4> : tensor<1x3xi32>  // Precalculated z1 * q2.

// CHECK: return %[[TRANSPOSE]]

// -----

// Tests that a float `stablehlo.transpose` is not converted to `tfl.transpose`.

func.func @transpose_float(%arg0: tensor<2x3x4xf32>) -> tensor<4x3x2xf32> {
  %0 = stablehlo.transpose %arg0, dims = [2, 1, 0] : (tensor<2x3x4xf32>) -> tensor<4x3x2xf32>
  return %0 : tensor<4x3x2xf32>
}
// CHECK-LABEL: transpose_float
// CHECK-NOT: tfl.transpose
// CHECK: stablehlo.transpose

// -----

func.func private @dense_image_warp_invalid_input_shape(%arg0: tensor<2x4x4xf32>, %arg1: tensor<2x4x4x2xf32>) -> tensor<2x4x4x1xf32> attributes {tf._implements = "addons:DenseImageWarp"}

// expected-warning @+1 {{Flow should be a 4D float tensor}}
func.func private @dense_image_warp_invalid_flow_shape(%arg0: tensor<2x4x4x1xf32>, %arg1: tensor<2x4x4xf32>) -> tensor<2x4x4x1xf32> attributes {tf._implements = "addons:DenseImageWarp"}