'dilation_h': 1,
        'padding': 'SAME',
        'act': 'RELU'
    }

    self._assertOpAndComposite([input_, filter_, bias],
                               tf.function(gen_mnist_ops.new_conv2d),
                               ops_defs._composite_conv_add_relu, kwargs)

  def test_new_conv2d_relu6(self):
    input_ = tf.random.uniform([1, 4, 4, 1])
    filter_ = tf.random.uniform([2, 2, 1, 8])

    # the same size as the input).

    # NOTE: The data/x/input is always specified in floating point precision.
    # output shape: [-1, 28, 28, 32]
    conv1 = gen_mnist_ops.new_conv2d(x, self.weights['f1'], self.biases['b1'],
                                     1, 1, 1, 1, 'SAME', 'RELU')

    # Max pooling. The kernel size spec {ksize} also follows the layout of

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