Currently, PlotCollection.wrap has a col_wrap parameter that defaults to 4 and needs to be an integer. We could allow it to also be None and then set it automatically like we used to do in arviz to get grids as close to sqrt(n) x sqrt(n) as possible.