Debugging flax models

[zaccharieramzi]

I would like to know what are the best practices to debug flax models.
I of course know of jax.debug.breakpoint and jax.debug.print but I would like to use VSCode's debugger in order to be more efficient.

I have noticed that even when I disable jitting jax.disable_jit(), I am unable to use VSCode's debugger or indeed any other form of manual debugging like print inside a @nn.compact call function.