Shape Parameter for Multiple Treatment Groups

[lucianopaz]

@jpryda

idx = np.arange(n_count_data)

You could change that to

idx = np.arange(n_count_data)[:, None]

This should allow the condition in the switch to broadcast well, and the indexing you do here:

@jpryda

mu=lambda_[count_data['person_id']]

should still work.