// tfl.Abs / tfl.Average_pool_2d / tfl.Cos / tfl.div / tfl.exp / tfl.hardswish /
// tfl.log / tfl.logistic / tfl.max_pool_2d / tfl.mirror_pad / tfl.maximum /
// tfl.custom / tfl.mean / tfl.minimum / tfl.pad / tfl.pow / tfl.prelu /
// tfl.relu / tfl.relu6 / tfl.rsqrt / tfl.sin / tfl.slice / tfl.softmax /
// tfl.space_to_depth / tfl.sqrt / tfl.square / tfl.squared_difference /
// tfl.strided_slice / tfl.tanh / tfl.transpose / tfl.transpose_conv

    func_op.eraseResult(0);
    func_op.insertResult(0, new_call_op.getResult(0).getType(),
                         /*resultAttrs=*/nullptr);

    // Modify the quantized fused function to do dequantize+relu(6).
    rewriter.setInsertionPoint(req_op);
    Value new_result = rewriter.create<mlir::stablehlo::UniformDequantizeOp>(
        req_op.getLoc(), func_op.getResultTypes()[0], req_op.getOperand());

  %3 = "tf.BiasAdd"(%2, %cst_0) {data_format = "NHWC", device = ""} : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32>
  %4 = "tf.Relu"(%3) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>
  %5 = "quantfork.qcast"(%4) : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2x!quant.uniform<i8:f32, 0.0027450981093387976:-19>>

      %outputs_12, %control_13 = tf_executor.island wraps "tf.Relu"(%outputs_10) {device = ""} : (tensor<16x112x112x?xf32>) -> tensor<16x112x112x?xf32>

  %broad2_s = "quantfork.stats"(%broad2) {layerStats = dense<[0.000000e+00, 1.000000e+01]> : tensor<2xf32>} : (tensor<?x26x26x26x16xf32>) -> tensor<?x26x26x26x16xf32>
  %add = "tfl.add"(%broad1_s, %broad2_s) {fused_activation_function = "RELU"} : (tensor<?x26x26x26x16xf32>, tensor<?x26x26x26x16xf32>) -> tensor<?x26x26x26x16xf32>

                        (TFL_RangeOp $start, $limit, $delta)>;
def LegalizeRelu6 : Pat<(TF_Relu6Op $arg), (TFL_Relu6Op $arg)>;
def LegalizeRelu : Pat<(TF_ReluOp $arg), (TFL_ReluOp $arg)>;
// TFL Relu doesn't support I32/I64 type, so legalizes TF Relu to TFL Maximum.
def LegalizeReluI32 :
  Pat<(TF_ReluOp TensorOf<[I32]>:$arg),
      (TFL_MaximumOp $arg,
        (Arith_ConstantOp ConstantAttr<RankedI32ElementsAttr<[]>,"0">))>;

  // CHECK-NEXT: tf_mlrt.executeop.device{{.*}}op: \22Relu\22
  // CHECK-NEXT: tf_mlrt_tpu.compile_and_execute
  %0 = "tf.Cast"(%arg0) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
  %1 = "tf.Relu"(%0) {__op_key = 1: i32, device = "/device:CPU:0"} : (tensor<f32>) -> (tensor<f32>)

        (NamedAttr<"equation"> $equation))),
    [(IsNotInLiftedFunc $res),
     (IsEinsumSupportedByXlaDotV2 $equation)],
    [], (addBenefit 10)>;

}
defm : LiftCompositeOpsWithActivation<TF_ReluOp, "relu">;
defm : LiftCompositeOpsWithActivation<TF_Relu6Op, "relu6">;

def LiftGather : Pat<
  (TF_GatherV2Op:$res $params, $indices, $axis, $batch_dims),
  (LiftAsTFPartitionedCall<"composite_gather_fn">

// CHECK-LABEL: func @relu_qint8
func.func @relu_qint8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8> {
  // CHECK: %[[X:.*]] = "tf.Relu"(%arg0) : (tensor<1xi8>) -> tensor<1xi8>
  %0 = "tf.Relu"(%arg0) : (tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8>
  func.return %0: tensor<1x!tf_type.qint8>
}

// -----

    %10 = "tf.Cast"(%9) {Truncate = false, device = ""} : (tensor<1x3xi32>) -> tensor<1x3xf32>
    %11 = "tf.Mul"(%10, %cst) {device = ""} : (tensor<1x3xf32>, tensor<f32>) -> tensor<1x3xf32>
    %12 = "tf.Relu"(%11) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
    return %12 : tensor<1x3xf32>
  }
// CHECK-LABEL: func @matmul_with_relu