Not sure if this comment fits better here or to one of the function above, but I was thinking that maybe at the ssms level, we could be helped by allowing theta to come in as a 3D-array where the first dimension is n_replicas e.g..
If we allow that as an input from the get-go, it could simplify some of the other logic?
Originally posted by @AlexanderFengler in #210 (comment)