(self, untyped_storage, *, dtype_hint=None, device_hint=None)
| 638 | # |
| 639 | # If we had a FakeTensor, device_hint tells us what device should be |
| 640 | def storage(self, untyped_storage, *, dtype_hint=None, device_hint=None) -> str: |
| 641 | ws = StorageWeakRef(untyped_storage) |
| 642 | v = self.seen_storages.get(ws) |
| 643 | if v is not None: |
| 644 | return v |
| 645 | v = f"buf{next(self.storage_counter)}" |
| 646 | maybe_dtype_hint = "" |
| 647 | if _dtype_or_default(None) != _dtype_or_default(dtype_hint): |
| 648 | maybe_dtype_hint = f", dtype_hint={dtype_hint!r}" |
| 649 | # TODO: being optional on device is kind of pointless as the default |
| 650 | # is CPU but most repros we care about are CUDA |
| 651 | maybe_device = "" |
| 652 | device = untyped_storage.device |
| 653 | if device.type == "meta": |
| 654 | assert device_hint is not None |
| 655 | device = device_hint |
| 656 | if _device_or_default(None) != device: |
| 657 | maybe_device = f", device={device!r}" |
| 658 | nbytes = untyped_storage.nbytes() |
| 659 | storage_hash = None |
| 660 | if self.store is not None and untyped_storage.device.type != "meta": |
| 661 | storage_hash = self.store.write_storage(untyped_storage) |
| 662 | self._lines.append( |
| 663 | f"{v} = reader.storage({storage_hash!r}, {nbytes!r}{maybe_device}{maybe_dtype_hint})" |
| 664 | ) |
| 665 | self.seen_storages[ws] = v |
| 666 | return v |
| 667 | |
| 668 | def tensor(self, name, t) -> None: |
| 669 | storage = self.storage( |
no test coverage detected