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>> {

              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),