I'm currently wrangling with pytorch.compile (flex attention doesn't like bias terms - see issues 145869 and 144511 if curious). As much as I love speed and optimization, JIT (at least the pytorch flavor) currently has weird return types that break VS Code's intellisense, weird stack traces, limitations around printing, and random issues like the bias term, and limitations such as not supporting sparsity.
Speaking of pytorch JIT workflows - what's a nice way of having a flag to turn off compilation?
I'd like to know more about this!
For some higher level context, see https://pytorch.org/blog/flexattention/
I'm alway shocked when people in this line of work either take things at face-value or just lie by omission/imprecision. this is not python breakpoints - it's a wholeass other system they had to re-roll that emulates python breakpoints:
> Unlike pdb, you will not be able to step through the execution, but you are allowed to resume it.
you think it works with pycharm/vscode/any other tooling? spoiler alert: of course not.
so no in JAX you can not.
Some of the current implementations are either strict Python subsets or very nearly so, but, yes, DSLs are distinct languages, that's what the L stands for.
https://docs.python.org/3.13/whatsnew/3.13.html#an-experimen...