pytorch
7550ec16 - Add support for dictionary with torch object keys. (#103158)

Commit
1 year ago
Add support for dictionary with torch object keys. (#103158) Fixes: #101979 This PR adds support for dictionaries with torch object as keys in dynamo. The main problem was that, for example, the source built for `d[torch.float]` (`d` being a dictionary) was `ODictGetItemSource(GlobalSource('d'), index=torch.float)`. When `Source.name` method was called, we got `odict_getitem(G['d'], torch.float)`. Evaluating that string raised an error, since `torch` was only available in the global dictionary `G` as `G["torch"]`. Instead, this PR builds the source: `ODictGetItemSource(GlobalSource('d'), index=AttrSource(GlobalSource('torch'), 'float'))`. The to-be-evaluated string is correctly generated as: `odict_getitem(G['d'], G['torch'].float)`. Here's a minimal example that reproduces the error, before this PR: ```python import torch d = { torch.float16: torch.float32, } @torch.compile def f(): return torch.randn(3, dtype=d[torch.float16]) f() ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/103158 Approved by: https://github.com/mlazos
Author
Committer
Parents
Loading