Repository navigation
[BugFix] Fix symbolic_trace for one output and the _ sentinel - #2646
Open
YeonwooSung wants to merge 4 commits into
Open
YeonwooSung wants to merge 4 commits into
YeonwooSung wants to merge 4 commits into
Conversation
A TensorDictModule that writes a single key traces a graph whose forward returns one tensor. Zipping out_keys against that tensor iterates dim 0 and raises ValueError.
A traced module with one output returns that tensor, not a 1-tuple. Wrap a non-tuple result so zip pairs out_keys with outputs instead of walking dim 0. A tuple of tensors is left unchanged.
"_" marks an output that eager TensorDictModule does not write. The traced
module unravels that key to ("_",) and the forward check against "_" never
matches, so the traced tensordict contains "_".
symbolic_trace stores out keys with _unravel_key_to_tuple, so "_" becomes
("_",). Compare against that tuple. A nested key ("_", "y") is not the
sentinel and is still written.
meta-cla
Bot
added
the
CLA Signed
github-actions
Bot
added
bug
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
tensordict.prototype.fxis a prototype.symbolic_tracereturns aTDGraphModulewhose forward is supposed to match the module that was traced. Two bugs in that forward made the traced result disagree with eagerTensorDictModule.A
TensorDictModulewith one out key traces a graph that returns a single tensor. The traced forward did_zip_strict(self.out_keys, outputs). A tensor iterates on dim 0, so zip raisedValueError: zip() argument 2 is longer than argument 1, or paired a batch row with the key when the lengths happened to match. The forward now wraps a non-tuple result in a 1-tuple. A graph that already returns a tuple of tensors is left unchanged."_"is the sentinel for "do not write this output". EagerTensorDictModuleskips it. The tracer stores out keys with_unravel_key_to_tuple, so"_"becomes("_",). The forward compared that tuple with the string"_", which is never equal, and wrote a key named"_". The check is nowout_key != ("_",). A real nested key("_", "y")is not the sentinel and is still written.Tests
Extended
test/nn/test_fx.py(no new test file):test_tensordictmodule_trace_single_out_keytracesTensorDictModule(nn.Linear(3, 2), in_keys=["x"], out_keys=["y"])and checks that keyymatches the eager values. Before the wrap, this raisedValueError: zip() argument 2 is longer than argument 1.test_tensordictmodule_trace_underscore_sentineltraces a module without_keys=["y", "_"]and checks that the traced keys equal the eager keys. Before the sentinel fix, the traced tensordict had an extra"_"key ({'_', 'y'}vs{'y'}).python -m pytest test/nn/test_fx.py -q --tb=short: 5 passed.Ruff check and format are clean on
tensordict/prototype/fx.pyandtest/nn/test_fx.py.