lhs_dilation_2d, rhs_dilation_2d, window_reversal_2d, dnums_2d,
        conv_op.getFeatureGroupCount(), conv_op.getBatchGroupCount(),
        conv_op.getPrecisionConfigAttr());

    OpResult conv2d_output = conv2d_op->getResult(0);
    auto conv2d_output_type = mlir::cast<ShapedType>(conv2d_output.getType());

    //
    // Transpose and reshape the output
    //