Expliclty avoid recording when should_record_events is false in record_shapeenv_event (#138965)

Looking at the function record_shapeenv_event its hard to tell that it does not always run
but we do disable it by setting top level is_recording to True self.should_record_events is false
this makes it more explicit to avoid confusion and overloading is_recording.

alternativley we can rename is_recording to do_no_record.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/138965
Approved by: https://github.com/ezyang
ghstack dependencies: #138804
This commit is contained in:
Laith Sakka 2024-10-25 19:20:44 -07:00 committed by PyTorch MergeBot
parent a688c57033
commit 475ba1df8d
2 changed files with 3 additions and 2 deletions

View File

@ -253,7 +253,8 @@ def record_shapeenv_event(*, save_tracked_fakes: bool = False) -> Callable:
return r
try:
if args[0].is_recording: # type: ignore[has-type]
shape_env = args[0]
if not shape_env.should_record_events or shape_env.is_recording: # type: ignore[has-type]
# If ShapeEnv is already recording an event, call the wrapped
# function directly.
#

View File

@ -2918,7 +2918,7 @@ class ShapeEnv:
)
# This will make sure we only record the top-level function call.
self.is_recording = not self.should_record_events
self.is_recording = False
# Keep track of the list of tracked fakes.
self.tracked_fakes = tracked_fakes
# List of events for reconstructing ShapeEnv at arbitrary points in time.