^bb0(%arg0: tensor<*xf32>, %arg1: tensor<5x3xi32>):
  %cst = arith.constant dense<2.0> : tensor<f32>
  // CHECK: "tfl.padv2"(%arg0, %arg1, %cst)
  %0 = "tfl.padv2"(%arg0, %arg1, %cst) : (tensor<*xf32>, tensor<5x3xi32>, tensor<f32>) -> tensor<? x f32>
  func.return %0#0 : tensor<? x f32>
}

// -----

// test PadV2 with invalid paddings size

// CHECK: %[[PAD:.+]] = "tfl.padv2"(%[[ARG0]], %[[PADDING]], %[[ARG1]]) : (tensor<2x3x!quant.uniform<i8:f32, 2.000000e+00:-1>>, tensor<2x2xi32>, tensor<!quant.uniform<i8:f32, 2.000000e+00:-1>>) -> tensor<4x5x!quant.uniform<i8:f32, 2.000000e+00:-1>>
// CHECK: return %[[PAD]]

// -----

// Tests that a quantized `stablehlo.pad` with interior padding is converted to
// `tfl.dilate` and `tfl.padv2`.

func.func @pad_with_interior_padding(

  // CHECK-LABEL: padv2
  // CHECK:  "tfl.padv2"(%arg0, %arg1, %cst) : (tensor<2x1x3xf32>, tensor<3x2xi32>, tensor<f32>) -> tensor<?xf32>
  // CHECK:  return
}

func.func @padv2_5D(tensor<2x1x3x1x1xf32>, tensor<5x2xi32>) -> tensor<? x f32> {
^bb0(%arg0: tensor<2x1x3x1x1xf32>, %arg1: tensor<5x2xi32>):
  %cst = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>

    ReplaceReturnOp(new_op.getElseBranch(), rewriter);
    rewriter.replaceOp(op, new_op.getResults());
    return success();
  }
};

// Converts mhlo.pad to tf.PadV2
Value ConvertPadOp(PatternRewriter& rewriter, Operation* old_op) {
  auto pad_op = cast<mhlo::PadOp>(old_op);
  mlir::Location loc = pad_op.getLoc();

        op, output_type, op.getOperands(), axis,
        /*fused_activation_function=*/rewriter.getStringAttr("NONE"));
  }
};

// Rewrites quantized stablehlo.pad to tfl.padv2.
// tfl.dilate is introduced in between when interior padding exists.
// TODO: b/322428814 - Add StableHLO quantizer integration tests for ODML.
class RewriteQuantizedPadOp : public OpRewritePattern<stablehlo::PadOp> {

  let results = (outs TFL_TensorOf<[F32, I32, I64, QI8, QUI8, TFL_Quint8, QI16]>:$output);

  let hasOptions = 1;

  let hasFolder = 1;
}

def TFL_PadV2Op : TFL_Op<"padv2", [
    PredOpTrait<"input and output must have same element type",
      TFL_TCresVTEtIsSameAsOp<0, 0>>,
    Pure,
    QuantizableResult,
    SameOperandsAndResultsScale,
    TFL_OperandHasRankAtMost<0, 5>,

  %paddings = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64>
  // expected-error @+1 {{failed to verify that operand 1 is 2-D}}
  %0 = "tf.PadV2"(%input, %paddings, %constant) : (tensor<2xi64>, tensor<2xi64>, tensor<i64>) -> tensor<3xi64>
  func.return %0 : tensor<3xi64>
}

// -----