Repository: scikit-learn/scikit-learn
File: sklearn/inspection/_plot/decision_boundary.py (Python)

Code before:
                    self.response.argmax(axis=2)
                    if self.response.ndim == 3
                    else self.response
                )
                self.surface_ = plot_func(
                    self.xx0, self.xx1, response, colors=colors, **kwargs
                )
            elif plot_method == "pcolormesh" and self.response.ndim == 2:  # predict
                cmap = mpl.colors.ListedColormap(colors)
                self.surface_ = plot_func(
                    self.xx0, self.xx1, self.response, cmap=cmap, **kwargs
                )
            else:
                multiclass_cmaps = [
                    mpl.colors.LinearSegmentedColormap.from_list(
                        f"colormap_{class_idx}",
                        [(1.0, 1.0, 1.0, 1.0), (r, g, b, 1.0)],
                    )
                    for class_idx, (r, g, b, _) in enumerate(colors)
                ]
                self.surface_ = []
                for class_idx, cmap in enumerate(multiclass_cmaps):
                    if self.response.ndim == 2:  # predict
                        response = np.ma.array(
                            (self.response == class_idx),
                            mask=~(self.response == class_idx),
                        )

                    else:  # predict_proba or decision_function
                        response = np.ma.array(
                            self.response[:, :, class_idx],
                            mask=~(self.response.argmax(axis=2) == class_idx),
                        )

                    self.surface_.append(
                        plot_func(self.xx0, self.xx1, response, cmap=cmap, **kwargs)
                    )

        if xlabel is not None or not ax.get_xlabel():
            xlabel = self.xlabel if xlabel is None else xlabel
            ax.set_xlabel(xlabel)
        if ylabel is not None or not ax.get_ylabel():
            ylabel = self.ylabel if ylabel is None else ylabel
            ax.set_ylabel(ylabel)

        self.ax_ = ax
        self.figure_ = ax.figure
        return self

    @classmethod
    def from_estimator(

Review comment (suggestion):
Is there a reason to prefer this one over the other? I find the latter a bit easier to parse.
```suggestion
                            mask=(self.response != class_idx),
```

Code after:
                    colors = plt.get_cmap("tab10", 10).colors[: self.n_classes]
                elif cmap == "tab20" and self.n_classes <= 20:
                    colors = plt.get_cmap("tab20", 20).colors[: self.n_classes]
                else:
                    cmap = plt.get_cmap(cmap, self.n_classes)
                    if not hasattr(cmap, "colors"):
                        # Get `LinearSegmentedColormap` for non-qualitative cmaps
                        colors = cmap(np.linspace(0, 1, self.n_classes))
                    else:
                        colors = cmap.colors
            elif isinstance(self.multiclass_colors, list):
                colors = [mpl.colors.to_rgba(color) for color in self.multiclass_colors]
            else:
                raise ValueError("'multiclass_colors' must be a list or a str.")

            self.multiclass_colors_ = colors

            if self.response.ndim == 2:  # predict
                # `pcolormesh` requires cmap, for the others it makes no difference
                cmap = mpl.colors.ListedColormap(colors)
                self.surface_ = plot_func(
                    self.xx0, self.xx1, self.response, cmap=cmap, **kwargs
                )

            # predict_proba and decision_function differ for plotting methods
            elif plot_method == "contour":
                # Plot only integer class values
                self.surface_ = plot_func(
                    self.xx0,
                    self.xx1,
                    self.response.argmax(axis=2),
                    colors=colors,
                    **kwargs,
                )
            else:
                multiclass_cmaps = [
                    mpl.colors.LinearSegmentedColormap.from_list(
                        f"colormap_{class_idx}",
                        [(1.0, 1.0, 1.0, 1.0), (r, g, b, 1.0)],
                    )
                    for class_idx, (r, g, b, _) in enumerate(colors)
                ]
                self.surface_ = []
                for class_idx, cmap in enumerate(multiclass_cmaps):
                    response = np.ma.array(
                        self.response[:, :, class_idx],
                        mask=(self.response.argmax(axis=2) != class_idx),
                    )
                    self.surface_.append(
                        plot_func(self.xx0, self.xx1, response, cmap=cmap, **kwargs)
                    )