rewriter.replaceAllUsesExcept(rhs, dq.getOutput(), dq);
  }
};

// Splits hybrid quantized `stablehlo.convolution` into `tfl.dequantize` and
// float `stablehlo.convolution` op. Weight tensor is transposed to match the
// filter tensor format for TFLite convolution.
// Legalization of float `stablehlo.convolution` op relies on existing passes
// for conversion of StableHLO -> MHLO -> TF -> TFL.

  return %0 : tensor<1x3x3x2x!quant.uniform<i8:f32, 4.000000e+0>>
}
// 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

// -----

            .getDimSize(input_feature_dimension);

    // Check for Group Convolution parameters
    if (feature_group_count != 1 && feature_group_count != input_channels) {
      // Group convolution is not supported yet.
      return rewriter.notifyMatchFailure(conv_op,
                                         "doesn't support group convolution");
    }

    auto input_spatial_dimensions = dnums.getInputSpatialDimensions();

      ]
  )
  def conv(self, conv_input: core.Tensor) -> Mapping[str, core.Tensor]:
    """Performs a 2D convolution operation.

    Args:
      conv_input: Input tensor to perform convolution on.

    Returns:
      A map of: output key -> output result.
    """
    out = nn_ops.conv2d(
        conv_input,
        self.conv_filters,

    QuantizableResult,
    TFL_SparseOp,
    DeclareOpInterfaceMethods<TFL_ArithmeticCount>,
    DynamicRangeQuantizedOpInterface]> {
  let summary = "Transpose convolution operator";

  let description = [{
    Performs transpose convolution operation on input.
  }];

  let arguments = (ins
    TFL_I32Tensor:$output_shape,
    TFL_TensorOf<[F32, QI8, QUI8, QI16]>:$weights,

      return mlir::cast<IntegerAttr>(attr).getInt() == 1;
    });
  };

  // Convolutions with 1x1 filter and with strides and dilations all ones, can
  // be computed as a GEMM in NHWC data format, and can be up to ~2x times
  // faster than convolution in NCHW.
  const bool one_by_one = d0 == 1 && d1 == 1;
  const bool trivial_strides = all_ones(getStrides());

  let summary = "Applies automatic space to depth transform for the first or frontier convolutions consume host inputs on TPU.";

  let description = [{
    Automatic space to depth transform is done by adding space to depth transform op after host input
    and applying space to depth transform for the first convolution and its backprop filter on TPU.

    For example, original program:

    ```mlir

      dnums.kernel_spatial_dimensions_size()) {
    return op.emitOpError() << "Both arguments to convolution must have "
                            << "same number of dimensions.\n";
  }

  if (dnums.input_spatial_dimensions_size() !=
      dnums.output_spatial_dimensions_size()) {
    return op.emitOpError() << "Both input and output of convolution must have "
                            << "same number of dimensions.\n";
  }