func.return %0 : tensor<?x?x?x?xf32>
}

// CHECK-LABEL: func @transposeConv2DBackpropFilter
func.func @transposeConv2DBackpropFilter(
  %input: tensor<1x32x32x3xf32>,
  %filter_sizes: tensor<4xi32>,
  %out_backprop: tensor<1x32x32x8xf32>
) -> tensor<1x1x3x8xf32> {

  // CHECK: %[[FILTER_PERM:[0-9]*]] = "tf.DataFormatVecPermute"
  // CHECK-SAME: dst_format = "NCHW"
  // CHECK-SAME: src_format = "NHWC"

  "$0.getType().cast<RankedTensorType>().getRank(), $1, &$_builder)">;

def LowerBiasAddGradOp :
  Pat<(TF_BiasAddGradOp AnyRankedTensor:$out_backprop, $data_format),
      (TF_SumOp $out_backprop,
                (TF_ConstOp (GetBiasAddGradReductionIndices $out_backprop,
                                                            $data_format)),
                /*keep_dims=*/ConstBoolAttrFalse)>;

def LegalizeMatrixDiag : Pat<(TF_MatrixDiagOp $diagonal),
                             (TFL_MatrixDiagOp $diagonal)>;

def LegalizeConv2DBackpropInput : Pat<
  (TF_Conv2DBackpropInputOp $input_sizes, $filter, $out_backprop,
     IsIntList1XY1:$strides,
     BoolAttr:$use_cudnn_on_gpu,
     IsSameOrValid:$padding,
     I64ArrayAttr:$explicit_paddings,
     IsDataFormatNHWC:$data_format,
     IsAllOnes:$dilations),

  func.return %0 : tensor<256x9x7x16xf32>
}

// -----

// CHECK-LABEL: @conv2d_backprop_input_dynamic
func.func @conv2d_backprop_input_dynamic(%filter: tensor<2x2x1x16xf32>, %out_backprop: tensor<?x256x256x16xf32>) -> tensor<?x512x512x1xf32> {
  // CHECK: %[[REV_FILTER:.*]] = "mhlo.reverse"(%arg0) <{dimensions = dense<[0, 1]> : tensor<2xi64>}>
  // CHECK: %[[RESULT:.*]] = mhlo.convolution(%arg1, %[[REV_FILTER]])

      return op.emitOpError(
          "requires out_backprop operand to have rank at least two with `NHWC` "
          "data format");
  } else {
    // Op definition requires data_format to be either NHWC or NCHW.
    DCHECK_EQ(format, tensorflow::TensorFormat::FORMAT_NCHW);
    if (!HasRankAtLeast(op.getOutBackprop(), 3))
      return op.emitOpError(
          "requires out_backprop operand to have rank at least three with "

  }];

  let description = [{
It accumulates all the values from out_backprop into the feature dimension.
For NHWC data format, the feature dimension is the last. For NCHW data format,
the feature dimension is the third-to-last.
  }];

  let arguments = (ins
    Arg<TF_NumberTensor, [{Any number of dimensions.}]>:$out_backprop,

    DefaultValuedOptionalAttr<TF_ConvnetDataFormatAttr, "\"NHWC\"">:$data_format

    auto filter_ty = mlir::dyn_cast<RankedTensorType>(op.getFilter().getType());

    // With the exception of out_backprop's batch dimension, out_backprop and
    // filter need to have static shape. Filter is validated here, out_backprop
    // is mostly validated at use.
    if (!out_backprop_ty || !filter_ty || !filter_ty.hasStaticShape())
      return failure();

  // expected-error @+1 {{requires out_backprop operand to have rank at least two with `NHWC` data format}}
  %0 = "tf.BiasAddGrad"(%arg0) {data_format = "NHWC"} : (tensor<3xf32>) -> tensor<3xf32>
  func.return %0 : tensor<3xf32>
}

// -----

func.func @testBiasAddGrad(%arg0: tensor<2x3xf32>) -> tensor<3xf32> {
  // expected-error @+1 {{requires out_backprop operand to have rank at least three with `NCHW` data format}}