[mccoyb]

Can you link dynamic shape support? Big if true — but I haven’t been able to find anything on it.

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)

[catgary]

It’s part of model export/serialization, it is documented here:

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#...

[mccoyb]

Thanks! See above — I don’t think this is exactly dynamic shape support.

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.

[catgary]

I think you’re right, so it is now as shape polymorphic as any framework that with an XLA backend can be.

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).

[mccoyb]

I can’t comment on the lowest leaf of this thread, but thanks for update! I’ll read through this section and see if my intuitions are wrong or right.