/*dimension_numbers=*/new_dimension_nums,
        /*feature_group_count=*/op.getFeatureGroupCountAttr(),
        /*batch_group_count=*/op.getBatchGroupCountAttr(),
        /*precision_config=*/op.getPrecisionConfigAttr());

    // Transpose the output of the `ConvolutionOp` back to the original op's
    // output shape so that users' shapes match.
    // [b, 0, 1, f] => [b, f, 0, 1]

        conv_op.getWindowReversalAttr(), conv_op.getDimensionNumbers(),
        conv_op.getFeatureGroupCount(), conv_op.getBatchGroupCount(),
        conv_op.getPrecisionConfigAttr());
    // For static case, replace the convolution op now.
    if (!is_dynamic_broadcast) {
      rewriter.replaceOp(mul_op, {new_conv});
    } else {