News for hackers. Stuff that matters.
Re:Micrograd.jl by catgary on Friday August 30, 2024 @11:27AM (#41401633)
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#...

Re:Micrograd.jl by mccoyb on Friday August 30, 2024 @11:31AM (#41401684)
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.

Re:Micrograd.jl by catgary on Friday August 30, 2024 @11:42AM (#41401771)
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).

Re:Micrograd.jl by mccoyb on Friday August 30, 2024 @11:36AM (#41401730)
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.