code: def plot_dropout_rate_heterogeneity(
      model,
      suptitle="Heterogeneity in Dropout Probability",
      xlabel="Dropout Probability p",
      ylabel="Density",
      suptitle_fontsize=14,
      **kwargs
  ):
      
      from matplotlib import pyplot as plt
  
      a, b = model._unload_params("a", "b")
      beta_mean = a / (a + b)
      beta_var = a * b / ((a + b) ** 2) / (a + b + 1)
  
      rv = stats.beta(a, b)
      lim = rv.ppf(0.99)
      x = np.linspace(0, lim, 100)
  
      fig, ax = plt.subplots(1)
      fig.suptitle(suptitle, fontsize=suptitle_fontsize, fontweight="bold")
  
      ax.set_title("mean: {:.3f}, var: {:.3f}".format(beta_mean, beta_var))
      ax.set_xlabel(xlabel)
      ax.set_ylabel(ylabel)
  
      fig.tight_layout(rect=[0, 0.03, 1, 0.95])
      plt.plot(x, rv.pdf(x), **kwargs)
      return ax
docs: Plot the estimated gamma distribution of p.
  
      p - (customers' probability of dropping out immediately after a transaction).
  
      Parameters
      ----------
      model: lifetimes model
          A fitted lifetimes model, for now only for BG/NBD
      suptitle: str, optional
          Figure suptitle
      xlabel: str, optional
          Figure xlabel
      ylabel: str, optional
          Figure ylabel
      kwargs
          Passed into the matplotlib.pyplot.plot command.
  
      Returns
      -------
      axes: matplotlib.AxesSubplot