Why can't I use a vscode debugger to debug jax code, specifically pure functions. I understand that they provide their own framework for debugging but vscode debugger is quite comfortable. Is this because vscode debugger in the first attempt captures compile time?

Why can't I use a vscode debugger to debug jax code, specifically pure functions. I understand that they provide their own framework for debugging but vscode debugger is quite comfortable. Is this because vscode debugger in the first attempt captures compile time?

• pytorch
• jax
• flax

Share Improve this question Follow asked 1 hour ago akshatakshat 2122 bronze badges New contributor akshat is a new contributor to this site. Take care in asking for clarification, commenting, and answering. Check out our Code of Conduct. 1

• Have you any sample project, toy project related to this idea, share to me? – Vy Do Commented 1 hour ago

Add a comment  | 

1 Answer 1

Sorted by: Reset to default 0

Why can't I use a vscode debugger to debug jax code, specifically pure functions.

You can use vscode interactive debugging – during tracing, which is when the Python part of the code is running. You can't (as far as I know) use vscode interactive debugging during execution of the compiled kernels that the traced code dispatches to.

For more information on JAX's computational model, you could start with JAX Key Concepts: Tracing and continue through the tutorials from there.

JAX provides some specific tools for debugging within this traced computational model; you can read about them at JAX: debugging.