def plot_covariate_groups(self, covariates, values, plot_baseline=True, **kwargs):
        """
        Produces a visual representation comparing the baseline survival curve of the model versus
        what happens when a covariate(s) is varied over values in a group. This is useful to compare
        subjects' survival as we vary covariate(s), all else being held equal. The baseline survival
        curve is equal to the predicted survival curve at all average values in the original dataset.

        Parameters
        ----------
        covariates: string or list
            a string (or list of strings) of the covariate(s) in the original dataset that we wish to vary.
        values: 1d or 2d iterable
            an iterable of the values we wish the covariate(s) to take on.
        plot_baseline: bool
            also display the baseline survival, defined as the survival at the mean of the original dataset.
        kwargs:
            pass in additional plotting commands.

        Returns
        -------
        ax: matplotlib axis, or list of axis'
            the matplotlib axis that be edited.


        Examples
        ---------
        >>> from lifelines import datasets, CoxPHFitter
        >>> rossi = datasets.load_rossi()
        >>> cph = CoxPHFitter().fit(rossi, 'week', 'arrest')
        >>> cph.plot_covariate_groups('prio', values=np.arange(0, 15), cmap='coolwarm')

        >>> # multiple variables at once
        >>> cph.plot_covariate_groups(['prio', 'paro'], values=[[0, 0], [5, 0], [10, 0], [0, 1], [5, 1], [10, 1]], cmap='coolwarm')

        >>> # if you have categorical variables, you can simply things:
        >>> cph.plot_covariate_groups(['dummy1', 'dummy2', 'dummy3'], values=np.eye(3))

        """
        from matplotlib import pyplot as plt

        covariates = _to_list(covariates)
        n_covariates = len(covariates)
        values = _to_array(values)
        if len(values.shape) == 1:
            values = values[None, :].T

        if n_covariates != values.shape[1]:
            raise ValueError("The number of covariates must equal to second dimension of the values array.")

        for covariate in covariates:
            if covariate not in self.hazards_.index:
                raise KeyError("covariate `%s` is not present in the original dataset" % covariate)

        set_kwargs_drawstyle(kwargs, "steps-post")

        if self.strata is None:
            axes = kwargs.pop("ax", None) or plt.figure().add_subplot(111)
            x_bar = self._norm_mean.to_frame().T
            X = pd.concat([x_bar] * values.shape[0])

            if np.array_equal(np.eye(n_covariates), values):
                X.index = ["%s=1" % c for c in covariates]
            else:
                X.index = [", ".join("%s=%g" % (c, v) for (c, v) in zip(covariates, row)) for row in values]
            for covariate, value in zip(covariates, values.T):
                X[covariate] = value

            self.predict_survival_function(X).plot(ax=axes, **kwargs)
            if plot_baseline:
                self.baseline_survival_.plot(ax=axes, ls=":", color="k", drawstyle="steps-post")

        else:
            axes = []
            for stratum, baseline_survival_ in self.baseline_survival_.iteritems():
                ax = plt.figure().add_subplot(1, 1, 1)
                x_bar = self._norm_mean.to_frame().T

                for name, value in zip(_to_list(self.strata), _to_tuple(stratum)):
                    x_bar[name] = value

                X = pd.concat([x_bar] * values.shape[0])
                if np.array_equal(np.eye(len(covariates)), values):
                    X.index = ["%s=1" % c for c in covariates]
                else:
                    X.index = [", ".join("%s=%g" % (c, v) for (c, v) in zip(covariates, row)) for row in values]
                for covariate, value in zip(covariates, values.T):
                    X[covariate] = value

                self.predict_survival_function(X).plot(ax=ax, **kwargs)
                if plot_baseline:
                    baseline_survival_.plot(
                        ax=ax, ls=":", label="stratum %s baseline survival" % str(stratum), drawstyle="steps-post"
                    )
                plt.legend()
                axes.append(ax)
        return axes