Home
last modified time | relevance | path

Searched refs:saved_tensors_hooks (Results 1 – 9 of 9) sorted by relevance

/aosp_15_r20/external/pytorch/torch/distributed/tensor/parallel/
H A Dinput_reshard.py42 cx: Optional[torch.autograd.graph.saved_tensors_hooks] = None
45 saved_tensor_hooks = torch.autograd.graph.saved_tensors_hooks(
/aosp_15_r20/external/pytorch/torch/autograd/
H A Dgraph.py208 class saved_tensors_hooks: class
287 class save_on_cpu(saved_tensors_hooks):
575 class _swap_with_cloned(saved_tensors_hooks):
/aosp_15_r20/external/pytorch/test/distributed/fsdp/
H A Dtest_checkpoint_wrapper.py356 orig_init = torch.autograd.graph.saved_tensors_hooks.__init__
357 torch.autograd.graph.saved_tensors_hooks.__init__ = patched_init
386 torch.autograd.graph.saved_tensors_hooks.__init__ = orig_init
/aosp_15_r20/external/pytorch/test/
H A Dtest_autograd.py568 ctx_1 = torch.autograd.graph.saved_tensors_hooks(lambda x: x, unpack)
569 ctx_2 = torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x)
593 ctx_1 = torch.autograd.graph.saved_tensors_hooks(lambda x: x, unpack)
1946 class RegisterPostNodeHook(torch.autograd.graph.saved_tensors_hooks):
9410 with torch.autograd.graph.saved_tensors_hooks(pack, lambda x: x):
9419 with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x):
9427 with torch.autograd.graph.saved_tensors_hooks(lambda x: 2 * x, lambda x: x / 2):
9435 with torch.autograd.graph.saved_tensors_hooks(lambda x: 2 * x, lambda x: x):
9456 with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x):
9464 with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x):
[all …]
/aosp_15_r20/external/pytorch/torch/utils/
H A Dcheckpoint.py1051 class _recomputation_hook(torch.autograd.graph.saved_tensors_hooks):
1099 class _checkpoint_hook(torch.autograd.graph.saved_tensors_hooks):
/aosp_15_r20/external/pytorch/docs/source/notes/
H A Dautograd.rst783 :class:`~torch.autograd.graph.saved_tensors_hooks` to register a pair of
808 with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
827 with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
836 with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x):
/aosp_15_r20/external/pytorch/docs/source/
H A Dautograd.rst353 .. autoclass:: torch.autograd.graph.saved_tensors_hooks
/aosp_15_r20/external/pytorch/test/inductor/
H A Dtest_compiled_autograd.py2660 with torch.autograd.graph.saved_tensors_hooks(
/aosp_15_r20/external/pytorch/test/dynamo/
H A Dtest_repros.py1104 with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):