5 ms·
It’s part of model export/serialization, it is documented here: https://jax.readthedocs.io/en/latest/export/export.html#support-for-reverse-mode-ad https://jax
by catgary 2y ago
It’s part of model export/serialization, it is documented here:
https://jax.readthedocs.io/en/latest/export/export.html#support-for-reverse-mode-ad 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#shape-poly https://jax.readthedocs.io/en/latest/export/shape_poly.html#...
- mccoyb 2y agoThanks! 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.
- deleted 2y ago[deleted]
- catgary 2y agoI 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 2y agoI 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.