Will jax.lax.while support dynamic shapes?

[femtomc]

How would parametric shapes work inside a construct like jax.lax.while? What are the restrictions? Is this even something that is intended to be supported?

[daskol]

Primitive while requires static shapes. As far as I understand, there are some complications with while for dynamic shapes it seems that there is not support for dynamic shapes on the low-level [1][2].