[catgary]

Yeah I was pretty enthusiastic about Julia for a year or two, even using it professionally. But honestly, JAX gives you (almost) everything Julia promises and its automatic differentiation is incredibly robust. As Python itself becomes a pretty reasonable language (the static typing improvements in 3.12, the promise of a JIT compiler) and JAX develops (it now has support for dynamic shape and AOT compilation) I can’t see why I’d ever go back.

The Julia repl is incredibly nice though, I do miss that.

[adgjlsfhk1]

IMO the python JIT support won't help very much. Python currently is ~50x slower than "fast" languages, so even if the JIT provides a 3x speedup, pure python will still be too slow for anything that needs performance. Sure it will help on the margins, but a JIT can't magically make python fast.

[catgary]

I’m only really thinking about ML/scientific computing workflows where all the heavy lifting happens in jax/torch/polars.

[adgjlsfhk1]

right and in those cases, a python JIT will do nothing for your performance because all the computation is happening in C/CUDA anyway.

[physicsguy]

It depends on whether you’re transforming data out of files or whatever to get it into these libraries to start with to be fair. Overall I wouldn’t expect that to have an effect on a long-running computation but when starting up a project it can be a bit slow.

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