conv_op) {
        return failure();
      }
      Value new_shape_of = rewriter.create<shape::ShapeOfOp>(
          mul_op.getLoc(), shape_of_op.getType(), new_conv);
      shape_of_op.replaceAllUsesWith(new_shape_of);
      rewriter.replaceOp(mul_op, {new_conv});
    }

    return success();
  }
};

class FuseMhloConvolutionPass
    : public PassWrapper<FuseMhloConvolutionPass, OperationPass<func::FuncOp>> {

    'Path to write the genereated register op file and MLIR file.')

flags.DEFINE_bool('gen_register_op', True,
                  'Generate register op cc file or tfr mlir file.')


@Composite(
    'NewConv2D',
    inputs=['input_: T', 'filter_: T', 'bias: T'],
    attrs=[
        'stride_w: int', 'stride_h: int', 'dilation_w: int', 'dilation_h: int',
        'padding: {"SAME", "VALID"}', 'act: {"", "RELU", "RELU6", "TANH"} = ""'

              RankedTensorType::get(
                  {static_cast<int64_t>(transposed_output_shape_i32.size())},
                  rewriter.getI32Type()),
              transposed_output_shape_i32));
      auto new_conv = rewriter.create<TF::Conv2DBackpropInputOp>(
          conv_op.getLoc(), output_type, output_sizes, filter, conv_input,
          rewriter.getI64ArrayAttr(strides),
          /*use_cudnn_on_gpu=*/rewriter.getBoolAttr(true),