Plan for IREE backend

[sebffischer]

Hi! I am interested in using IREE as a PjRt plugin, and I have seen that there used to be some activity in 2023 to make this work properly, but not so much since. I was wondering what the plans for this are? The big advantage I see is that the IREE compiler seems to support shape dynamism during runtime. Supporting IREE as a backend would (I believe) make jax better applicable in situations where we can basically expect every input to have a different shape.

[dishana11]

Hey! The IREE PJRt plugin for JAX is still in the works, maintained in IREE’s repo since its 2023 merge. It’s experimental but actively developed, also IREE’s dynamic shape support is a big win for handling varied input shapes in nlp,JAX etc. No big roadmap updates, but it’s moving forward. You can try it by building from source. check iree.dev/guides/ml-frameworks/jax/ for setup. Got any specific hardware in mind?

[sebffischer]

Thanks, this is great to hear! Is this being pushed from the jax side or is this up to the IREE team to do? I read on the IREE github that currently there is no one being paid to maintain the IREE PjRt plugin which I assumed is the way how jax would interface IREE, so this is why I am asking.