We could handle this case, assuming that `ax` is an array with size equal to `ngroups`. Want to open a separate issue for that?