"Add control flow callbacks for JIT-compiled frameworks" PEP

[MRiabov]

Hello Jax community,

I think if something like this PEP will be implemented in Python, Jax will have much easier time tracing what is happening in the code without cond and other terse code. A straightforward callback that would tell jit-compiled languages when if and other control flow keywords are happening.

PEP: https://discuss.python.org/t/add-control-flow-callbacks-for-jit-compiled-frameworks/75636/2

I've tagged @jakevdp in the PEP because he will probably know what is necessary for such callback.
Basically to replace cond with if.

I'd also think that if jax does a little smart work (internally inferences when to use vmap, or tree.map, or scan instead of for loops, which will be easily detectable now), most of jax "terseness" will be removed.

Cheers,
Maksym Riabov

[jakevdp]

Thanks for pointing out the discussion! In general I'm in agreement with the folks responding there: I don't think the proposed feature would be a net benefit to the Python language spec.

[MRiabov]

Actually it can be easily resolved via something like a C++ generic on python backend (I suppose C has them too). and would have no performance impact whatsoever. Do I understand correctly that C generics have no overhead except memory and compile time increase?

…

--- Well, anyway, when Mojo language will be develop the GPU computation well, python slowness will stop being an issue

On Tue, Dec 31, 2024 at 4:26 PM Jake Vanderplas ***@***.***> wrote: Thanks for pointing out the discussion! In general I'm in agreement with the folks responding there: I don't think the proposed feature would be a net benefit to the Python language spec. — Reply to this email directly, view it on GitHub <#25691 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AZZOTD32YJYIYVRDWDPPYPL2IKSQPAVCNFSM6AAAAABUKXALGSVHI2DSMVQWIX3LMV43URDJONRXK43TNFXW4Q3PNVWWK3TUHMYTCNZQGU4TIOA> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[jakevdp]

I'm not saying "it can't be done", I'm saying "doing it would not be a net benefit for the language". Added complexity in a language always carries a cost – you're proposing to add overloading to fundamental control flow that will make simple programs more difficult to reason about. I don't think that would be a good idea.