Starting from a setup very similar to yours

if one increases the height of the box while keeping everything else constant, much of the vortex structure is lost

Therefore, I think that the structure we can see in the example mostly comes down to the periodic boundaries (implicit in the array rolling) rather than classical vortex shedding behind a cylinder.
In any case, thank you very much for putting the script up, especially regarding handling the obstacle it proved very useful when I wrote my own small LBM simulator in JAX :)
Sadly, for standard setups, namely https://wwwold.mathematik.tu-dortmund.de/lsiii/cms/papers/SchaeferTurek1996.pdf, I think one would need to handle the boundaries properly, which will take some extra effort.
Anyway, since you're also doing lots of stuff in JAX, maybe you're interested in my implementation: https://github.com/leo1200/micro-lbm :D - and if you know a setup which works better, please let me know :)