[direct-linearize] fix name_stack bugs
Surprisingly, the bug was tracked down to #26111 aka cl/730939406, specifically
the new implementation of reset_name_stack in source_info_util.py.
To repro, use the before-this-commit implementation of reset_name_stack (left
commented-out in the file), and run
```
JAX_USE_DIRECT_LINEARIZE=1 python tests/name_stack_test.py NameStackTransformationTest.test_nested_jit_stack
```