Dynamic shaping, "round" function, JAX, and a "few" more questions

Thank you @ricardoV94 for your reply. I have just created a separate thread about this error I got. (Sampler initialization error with a model containing an LKJCholeskyCov distribution)