namespace tensorflow {

#define REGISTER_XLA_OPS_ON_DEVICE(DEVICE)                                     \
  REGISTER_KERNEL_BUILDER(Name("XlaConv")                                      \
                              .HostMemory("window_strides")                    \
                              .HostMemory("padding")                           \
                              .HostMemory("lhs_dilation")                      \

  (StableHLO_ConvolutionOp:$res $lhs, $rhs, $window_strides, $padding,
      $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
      $feature_group_count, $batch_group_count, $precision_config),
  (LiftAsTFXlaCallModule<"composite_conv_fn">
    (ArgumentList $lhs, $rhs),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),

    (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding,
        $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
        $feature_group_count, $batch_group_count, $precision_config),
    $bias),
  (LiftAsTFXlaCallModule<"composite_conv_with_bias_same_shape_fn">
    (ArgumentList $lhs, $rhs, $bias),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),

// CHECK: <{window_dimensions = array<i64: 1, 2, 2, 1>, window_strides = array<i64: 1, 2, 2, 1>}>
// CHECK: ^bb0(%[[REDUCE_ARG_0:.+]]: tensor<f32>, %[[REDUCE_ARG_1:.+]]: tensor<f32>):
// CHECK: %[[MAX:.+]] = stablehlo.maximum %[[REDUCE_ARG_0]], %[[REDUCE_ARG_1]]
// CHECK: stablehlo.return %[[MAX]]

// Check that the attributes window_dimensions & window_strides are also
// permutated to match the new input shape.

// CHECK:           %[[REDUCE:.*]] = "stablehlo.reduce_window"(%[[INPUT]], %[[PAD_VAL]])
// CHECK:               <{window_dimensions = array<i64: 1, 3, 3, 1>,
// CHECK:                 window_strides = array<i64: 1, 2, 2, 1>}> ({
// CHECK:           ^bb0(%[[ARG0:.*]]: tensor<f32>, %[[ARG1:.*]]: tensor<f32>):
// CHECK:             %[[MAX:.*]] = stablehlo.maximum %[[ARG0]], %[[ARG1]] : tensor<f32>

                           %weight_scale, %weight_zp) {
        Tlhs = "tfdtype$DT_FLOAT",
        Trhs = "tfdtype$DT_QINT8",
        Tout = "tfdtype$DT_FLOAT",
        window_strides = [1, 1],
        padding = "",
        explicit_padding = [],
        lhs_dilation = [],
        rhs_dilation = [],
        dimension_numbers = "",
        batch_group_count = 1,
        feature_group_count = 1,

  }
  if (const auto* op = op_union.AsStablehloConvolutionOptions()) {
    if (!(op->window_strides.empty())) {
      std::vector<int64_t> shape;
      shape.push_back(static_cast<int64_t>(op->window_strides.size()));
      attributes.emplace_back(builder.getNamedAttr(
          "window_strides",
          BuildVhloTensorV1Attr(shape, op->window_strides, builder)));
    } else {

      // $broadcast_dims
      CompileTimeConstantOperand<TF::XlaBroadcastHelperOp, 2>,
      // $window_strides, $padding, $lhs_dilation, $rhs_dilation,
      // $feature_group_count
      CompileTimeConstantOperand<TF::XlaConvOp, 2, 3, 4, 5, 6>,
      // $window_strides, $padding, $lhs_dilation, $rhs_dilation,
      // $feature_group_count
      CompileTimeConstantOperand<TF::XlaConvV2Op, 2, 3, 4, 5, 6>,

  if (op_name == "stablehlo.reduce_window" &&
      (field_name == "window_dimensions" || field_name == "window_strides" ||
       field_name == "base_dilations" || field_name == "window_dilations"))
    return true;
  if (op_name == "stablehlo.select_and_scatter" &&
      (field_name == "window_dimensions" || field_name == "window_strides"))
    return true;
  if (op_name == "stablehlo.slice" &&

        mlir::cast<TensorType>(op->getResult(0).getType());
    const TensorType new_conv_output_tensor_type =
        GetTransposedTensorType(output_tensor_type, kNchwToNhwcPermutation);

    // window_strides, padding, lhs_dilation, rhs_dilation, window_reversal are
    // reused without modification because the ordering of spatial dimensions
    // is not modified (i.e. before: [b, f, 0, 1], after: [b, 0, 1, f] => the