5 ms·

I hope this is the future for Python. Write in pure Python, but if needed, the code can be JIT (or not JIT) compiled into something faster (provided your code d

I hope this is the future for Python. Write in pure Python, but if needed, the code can be JIT (or not JIT) compiled into something faster (provided your code does not rely too much on low-level python stuff, such as __ functions).

Scene_Cast2 2y ago
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?
- dleeftink 2y ago
  > random issues like the bias term I'd like to know more about this!
  - Scene_Cast2 2y ago
    Basically this style of code - https://github.com/pytorch-labs/attention-gym/pull/84/files 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/pull/137452 https://github.com/pytorch/pytorch/issues/144511 https://github.com/pytorch/pytorch/issues/144511 https://github.com/pytorch/pytorch/issues/145869 https://github.com/pytorch/pytorch/issues/145869 For some higher level context, see https://pytorch.org/blog/flexattention/ https://pytorch.org/blog/flexattention/
- singhrac 2y ago
  You can decorate a function with @torch.compiler.disable() to disable a specific function (and anything further down the call stack).
dec0dedab0de 2y ago
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.
- O5vYtytb 2y ago
  There already is. https://docs.python.org/3.13/whatsnew/3.13.html#an-experimental-just-in-time-jit-compiler https://docs.python.org/3.13/whatsnew/3.13.html#an-experimen...
- masklinn 2y ago
  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.
sevensor 2y ago
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.
- almostgotcaught 2y ago
  worse is you can't set a breakpoint inside the "jitted" function (maybe can't print either...)
  - cl3misch 2y ago
    In JAX you can: https://jax.readthedocs.io/en/latest/debugging/print_breakpoint.html#interactive-inspection-with-jax-debug-breakpoint https://jax.readthedocs.io/en/latest/debugging/print_breakpo...
    - almostgotcaught 2y ago
      > 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.
  - ByteMe95 2y ago
    You can in CSP nodes https://github.com/Point72/csp https://github.com/Point72/csp
- dragonwriter 2y ago
  > 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.

Scene_Cast2 2y ago

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?

dleeftink 2y ago
> random issues like the bias term I'd like to know more about this!
- Scene_Cast2 2y ago
  Basically this style of code - https://github.com/pytorch-labs/attention-gym/pull/84/files 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/pull/137452 https://github.com/pytorch/pytorch/issues/144511 https://github.com/pytorch/pytorch/issues/144511 https://github.com/pytorch/pytorch/issues/145869 https://github.com/pytorch/pytorch/issues/145869 For some higher level context, see https://pytorch.org/blog/flexattention/ https://pytorch.org/blog/flexattention/
singhrac 2y ago
You can decorate a function with @torch.compiler.disable() to disable a specific function (and anything further down the call stack).

dec0dedab0de 2y ago

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.

O5vYtytb 2y ago
There already is. https://docs.python.org/3.13/whatsnew/3.13.html#an-experimental-just-in-time-jit-compiler https://docs.python.org/3.13/whatsnew/3.13.html#an-experimen...
masklinn 2y ago
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.

sevensor 2y ago

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.

almostgotcaught 2y ago
worse is you can't set a breakpoint inside the "jitted" function (maybe can't print either...)
- cl3misch 2y ago
  In JAX you can: https://jax.readthedocs.io/en/latest/debugging/print_breakpoint.html#interactive-inspection-with-jax-debug-breakpoint https://jax.readthedocs.io/en/latest/debugging/print_breakpo...
  - almostgotcaught 2y ago
    > 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.
- ByteMe95 2y ago
  You can in CSP nodes https://github.com/Point72/csp https://github.com/Point72/csp
dragonwriter 2y ago
> 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.