Edit: I see — I think you mean exporting lowered StableHLO code in a shape polymorphic format —- from the docs: https://jax.readthedocs.io/en/latest/export/shape_poly.html

This is not the thing I usually think when someone says dynamic shape support.

In this model, you have to construct a static graph initially —- then you’re allowed to specify a restricted set of input shapes to be symbolic, to avoid the cost of lowering — but you’ll still incur the cost of compilation for any new shapes which the graph hasn’t been specialized for (because those shapes affect the array memory layouts, which XLA needs to know to be aggressive)

https://jax.readthedocs.io/en/latest/export/export.html#supp...

Edit: I think you need to look here as well, the Exported objects do in fact serialize a function and support shape polymorphism:

https://jax.readthedocs.io/en/latest/export/shape_poly.html#...

My definition might be wrong, but I often think of full dynamic shape support as implying something dynamic about the computation graph.

For instance, JAX supports a scan primitive — whose length must be statically known. With full dynamic shape support, this length might be unknown — which would mean one could express loops with shape dependent size.

As far as I can tell, shape polymorphic exports may sort of give you that — but you incur the cost of compilation, which will not be negligible with XLA.

I work with edge devices, so I have also been experimenting with IREE for deployment, that can handle dynamic shapes (at times, it stopped working for a version but may be working again in the development branch).