jax
fe26c19b - [direct-linearize] fix name_stack bugs

Commit
336 days ago
[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 ```
Author
Committer
Parents
Loading