Just a quick note, it looks like some of the arguments for diagnostics.plot_sbc_histograms() are not working in the current version. Customizing n_row and n_col has no effect on the output, and providing them to trainer.diagnose_sbc_histograms via plot_args (e.g. plot_args = {"n_row": 3, "n_col": 4}) gives the following error:
TypeError: plot_sbc_histograms() got an unexpected keyword argument 'n_row'