All of that would be possible with the changes parent poster proposes. See the gotchas section in JAX, which is exactly these limitations:

https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_J...