Home
last modified time | relevance | path

Searched refs:WeakTensorKeyDictionary (Results 1 – 7 of 7) sorted by relevance

/aosp_15_r20/external/pytorch/torch/distributed/optim/
H A Dapply_optimizer_in_backward.py12 param_to_optim_hook_handle_map = torch.utils.weak.WeakTensorKeyDictionary()
13 param_to_acc_grad_map = torch.utils.weak.WeakTensorKeyDictionary()
/aosp_15_r20/external/pytorch/torch/testing/_internal/
H A Dlogging_tensor.py10 from torch.utils.weak import WeakTensorKeyDictionary
106 self.memo = WeakTensorKeyDictionary()
/aosp_15_r20/external/pytorch/torch/nested/_internal/
H A Dnested_tensor.py8 from torch.utils.weak import WeakTensorKeyDictionary
12 _tensor_symint_registry = WeakTensorKeyDictionary()
/aosp_15_r20/external/pytorch/torch/
H A D_guards.py32 from torch.utils.weak import WeakTensorKeyDictionary
664 self.tensor_to_context = WeakTensorKeyDictionary()
/aosp_15_r20/external/pytorch/torch/utils/
H A Dweak.py303 WeakTensorKeyDictionary = WeakIdKeyDictionary variable
/aosp_15_r20/external/pytorch/torch/fx/experimental/
H A Dproxy_tensor.py73 from torch.utils.weak import _WeakHashRef, WeakIdKeyDictionary, WeakTensorKeyDictionary
1015 self.tensor_tracker = WeakTensorKeyDictionary()
1398 self.tensor_tracker = WeakTensorKeyDictionary()
/aosp_15_r20/external/pytorch/test/inductor/
H A Dtest_torchinductor.py93 from torch.utils.weak import WeakTensorKeyDictionary
10742 live_tensors = WeakTensorKeyDictionary()