If this is the way forward, I'd love for a better developer experience.

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?

> random issues like the bias term

I'd like to know more about this!

Basically this style of code - https://github.com/pytorch-labs/attention-gym/pull/84/files - has issues like this - https://github.com/pytorch/pytorch/pull/137452 https://github.com/pytorch/pytorch/issues/144511 https://github.com/pytorch/pytorch/issues/145869

For some higher level context, see https://pytorch.org/blog/flexattention/

You can decorate a function with @torch.compiler.disable() to disable a specific function (and anything further down the call stack).

I would rather a JIT just built into the reference implementation. a JIT would help way more programs than removing the GIL but everyone thinks the GIL affects them for some reason.

They serve different use cases. The function JIT pattern being a manual opt in it can be much more aggressive by only supporting restricted language patterns rather than... the entire thing. They can also use bespoke annotations for better codegen e.g. you can tell numba to only codegen for i32 -> i32 -> i32, rather than lazily codegen for any a -> b -> c.

There already is.

https://docs.python.org/3.13/whatsnew/3.13.html#an-experimen...

The misleading thing about this approach is that the decorated function is no longer Python at all. It’s another language with Python syntax. Which is a neat way to get a parser for free, but it’s going to set up expectations about semantics that are bound to be incorrect.

> The misleading thing about this approach is that the decorated function is no longer Python at all. It’s another language with Python syntax.

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.

worse is you can't set a breakpoint inside the "jitted" function (maybe can't print either...)

In JAX you can: https://jax.readthedocs.io/en/latest/debugging/print_breakpo...

> In JAX you can

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.

You can in CSP nodes https://github.com/Point72/csp

It's funny how close it is to the title of my talk "Python already has a frontend for your compiler": https://github.com/true-grue/python-dsls