You can use dims, the problem is that due to the linked bug, shape is not being properly inferred by them so you need to be redundant and specify shape as well.
Code will benefit from vectorization if you use a single likelihood, so I would advise it.