possible path to shape dynamism

[sebffischer]

I think IREE supports shape dynamism during runtime: https://iree.dev/developers/usage-best-practices/ and there is an experimental pjrt plugin for it: https://github.com/iree-org/iree/tree/main/integrations/pjrt

@dfalbel

[sebffischer]

So it looks like the pjrt plugin itself is not super stable: iree-org/iree#19459, but the good thing is that IREE supports stableHLO, so we would only have to replace the {pjrt} package with an {iree} package to support shape dynamism I think. Not sure how easy this would be, though.

As a first step I should just check how well the iree plugins work. Unfortunately the latest supported PjRt version is 0.68 I think.

[sebffischer]

I asked about the plans of the jax team on the IREE backend here: jax-ml/jax#31611