Dynamic shape roadmap?

[thisiscam]

Hi JAX team,

I have seen sporadically in the git issues/PRs (e.g. #12497) that there were and probably will be efforts in bringing dynamic shape support to JAX. Further, currently there seems to be some minimal support if I enable it by export JAX_DYNAMIC_SHAPES=1.

Out of curiosity, could the team shed some light on the general roadmap for djax? More specifically:

• Will there be a timeline?
• Will jit be supported eventually for dynamic shapes?

[mattjj]

Thanks for the question, @thisiscam.

Will there be a timeline?

We are not that organized!

We have a technical plan, but we're making slow and unsteady progress on it as other more urgent stuff keeps arising.

Will jit be supported eventually for dynamic shapes?

Yes. Also JAX's other APIs, like vmap and grad, and those aren't trivial: if you vmap a dynamically-shaped function (even just vmap(jnp.arange), you might get ragged intermediates and outputs.

[alonfnt]

Hey @mattjj just wondering if there's any update on this topic.