// CHECK: %[[q:.*]] = "tfl.quantize"(%[[cst]]) <{qtype = tensor<1x1x3x!quant.uniform<i8<-127:127>:f32, 0.028479868971456691>>}> {volatile} : (tensor<1x1x3xf32>) -> tensor<1x1x3x!quant.uniform<i8<-127:127>:f32, 0.028479868971456691>>
// CHECK: %[[dq:.*]] = "tfl.dequantize"(%[[q]]) : (tensor<1x1x3x!quant.uniform<i8<-127:127>:f32, 0.028479868971456691>>) -> tensor<1x1x3xf32>

//
// This is needed to support Conv+BatchNorm pattern from Jax models converted
// using jax2tf w/o native serialization. Note that Jax2tf patterns always
// extend bias shapes to a rank of 4, e.g. 1x1x1x5.
def ConvertSubMulToMulAdd : Pat<
  (TF_MulOp
    (TF_SubOp
      (SupportedAffineOpMatcher $conv_out, $input, $weight),
      (TF_ConstOp:$sub_rhs IsFloatElementsAttr:$sub_rhs_value)),

  %cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
  %0 = "tf.XlaCallModule"(%arg0, %cst) {
    Sout = [#tf_type.shape<1x2x2x2>], _entry_function = @composite_conv_fn,
    _original_entry_function = "composite_conv_fn",
    _stablehlo_module_attrs = {}, _quantization_method = "weight_only_ptq { }",
    device = "", dim_args_spec = [], disabled_checks = [],

// CHECK: %[[DEQUANT:.*]] = stablehlo.uniform_dequantize %[[CONV]] : (tensor<1x4x3x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x4x3xf32>
// CHECK: return %[[DEQUANT]] : tensor<1x4x3xf32>

func.func @testFoldStridedSliceShapeWithBeginMask(%arg0: tensor<1x2x3x?xf32>) -> (tensor<2xi32>) {
  %0 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
  %1 = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
  %2 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
  %3 = "tf.Shape"(%arg0) : (tensor<1x2x3x?xf32>) -> tensor<4xi32>

// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<2x2x3x!quant.uniform<i8:f32, {{.*}}>) -> tensor<2x2x3xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<2x2x3xf32>

}

// CHECK-LABEL: @DontReorderReshapex2Add
func.func @DontReorderReshapex2Add(%arg0: tensor<1x2x3x4x!quant.uniform<i8:f32, 3.0>>, %arg1: tensor<1x2x3x4x!quant.uniform<i8:f32, 5.0>>) -> tensor<6x4x!quant.uniform<i8:f32, 7.0>> {
  %shape = arith.constant dense<[6, 4]> : tensor<2xi32>
  %0 = "tfl.reshape"(%arg0, %shape) : (tensor<1x2x3x4x!quant.uniform<i8:f32, 3.0>>, tensor<2xi32>) -> tensor<6x4x!quant.uniform<i8:f32, 3.0>>

  %0 = "tf.StridedSlice"(%arg0, %arg1, %arg2, %arg3) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64, offset = false} : (tensor<12x2x2x5x!tf_type.string>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1x2x2x5x!tf_type.string>
  func.return %0 : tensor<1x2x2x5x!tf_type.string>

// -----

// CHECK-LABEL: func @sign_dynamic
func.func @sign_dynamic(%arg0: tensor<?x2x3x?xf32>) -> tensor<?x2x3x?xf32> {
  // CHECK: [[SIGN:%.*]] = mhlo.sign %arg0 : tensor<?x2x3x?xf32>
  // CHECK: return [[SIGN]] : tensor<?x2x3x?xf32>
  %0 = "tf.Sign"(%arg0) : (tensor<?x2x3x?xf32>) -> (tensor<?x2x3x?xf32>)
  func.return %0 : tensor<?x2x3x?xf32>
}

// -----

// CHECK-LABEL: slice_constant_start

  func.return %1 : tensor<1x2x2x5xf32>

// CHECK: %0 = "tfl.dequantize"(%arg0)
// CHECK: %1 = "tfl.strided_slice"(%0, %arg1, %arg2, %arg3)
// CHECK: %2 = "tfl.quantize"(%1) <{qtype = tensor<1x2x2x5x!quant.uniform<u8:f32, 1.000000e-01>>}> {volatile}
// CHECK: %3 = "tfl.dequantize"(%2)
// CHECK: return %3 : tensor<1x2x2x5xf32>
}

// CHECK-LABEL: QuantizePad