Shape Parameter for Multiple Treatment Groups

TIL that np.arange(n_count_data)[:, None] works! This looks to be equivalent to np.expand_dims(np.arange(n_count_data), -1) (both turn the index into a n_count_data x 1 matrix instead of an array).