// CHECK-LABEL: testConv3dWithFloatInput
func.func @testConv3dWithFloatInput(%arg0: tensor<?x?x?x?x?xf32>,%arg1:  tensor<?x?x?x?x?xf32>,%arg2: tensor<?xf32>) -> tensor<?x?x?x?x?xf32> {
  // CHECK: "tfl.conv_3d"(%arg0, %arg1, %arg2)

  %0 = "tf.Conv3D"(%arg0, %arg1) {padding = "SAME", strides = [1, 1, 1, 1, 1]} : (tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
  func.return %0: tensor<?x?x?x?x?xf32>

  // CHECK-LABEL: conv3d_valid
  // CHECK:  %[[CST:.*]] = "tfl.no_value"() <{value}> : () -> none

}
// Confirm that the `stablehlo.convolution` is not converted to `tfl.conv_2d`.
// CHECK-LABEL: convolution_upstream_srq_non_const_filter
// CHECK-SAME: %[[ARG:.+]]: tensor<1x3x3x4x!quant.uniform<i8:f32, 1.000000e+00:-100>>
// CHECK: stablehlo.convolution
// CHECK-NOT: tfl.conv_2d

// -----

// Tests that if the window padding contains values of 0, tfl.pad op is not

      auto conv_result = rewriter.create<mhlo::ConvolutionOp>(
          conv_op.getLoc(), new_output_type, sliced_input, sliced_kernel,
          conv_op.getWindowStridesAttr(), conv_op.getPaddingAttr(),
          conv_op.getLhsDilationAttr(), conv_op.getRhsDilationAttr(),
          conv_op.getWindowReversalAttr(), conv_op.getDimensionNumbers(), 1, 1,
          conv_op.getPrecisionConfigAttr());

      conv_results.push_back(conv_result);

    TFL_ResourceTensor:$resource_id
  );

  let results = (outs TFL_TensorOf<[F32, F64, I1, UI8, I8, QI8, QUI8, I32, I64, QI16, Complex<F<32>>, Complex<F<64>>]>:$result);
}

def TFL_Conv3DOp : TFL_Op<"conv_3d", [
    Pure,
    AccumulatorUniformScale<2, 0, 1>,
    TFL_OperandHasRank<0, 5>,
    TFL_OperandHasRank<1, 5>,
    // Channel dimension in input and filter should match.

          ]
      )
      def conv3d(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
        """Performs a 3D convolution operation.

        Args:
          input_tensor: Input tensor to perform convolution on.

        Returns:
          A map of: output key -> output result.
        """
        out = nn_ops.conv3d(
            input_tensor,
            self.filters,

    func.return %arg0 : tensor<*xi32>
  }

  // Test conv2d inferReturnTypes can infer some information when input or
  // filter does not have fully static shape.

  // CHECK-LABEL: func @conv2d_unranked_input_and_filter
  func.func @conv2d_unranked_input_and_filter(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
    // CHECK: "tf.Conv2D"
    // CHECK-SAME: -> tensor<?x?x?x?xf32>

  %0 = "tf.Conv2D"(%arg0, %arg1) {padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<*xf32>, tensor<*xf32>) -> tensor<?x?x?x?xf32>
  func.return %0 : tensor<?x?x?x?xf32>
}

// -----

// CHECK-LABEL: func @testValidConv3D
func.func @testValidConv3D(%arg0: tensor<256x32x32x32x3xf32>, %arg1: tensor<3x3x3x3x16xf32>) -> tensor<256x32x32x32x16xf32> {

            FindUserOfType<TFL::QuantizeOp>(op) != nullptr);
  }
};

// Rewrites `stablehlo.convolution` into fused `tfl.conv_2d`.
// If available, fuse bias and activation adjacent to `stablehlo.convolution`.
// This RewritePattern rewrites both the following into `tfl.conv_2d` op:
//
// StableHLO Quantizer output:
//   * input: per-tensor qi8
//   * filter: per-channel qi8 (`quantization_dimension` = 3)

// with a 0-d constant, e.g. before this optimization,
//   %cst = arith.constant dense<1.0> : tensor<16x16x4xf32>
//   %0 = "tfl.conv_2d"...
//   %1 = "tfl.add"(%0, %cst) : (tensor<16x16x4xf32>, tensor<16x16x4xf32>)
// After this optimization:
//   %cst = arith.constant dense<1.0> : tensor<f32>
//   %0 = "tfl.conv_2d"...
//   %1 = "tfl.add"(%0, %cst) : (tensor<16x16x4xf32>, tensor<f32>)