1# mypy: allow-untyped-defs 2from typing import Any 3 4import torch 5 6# This module is defined in torch/csrc/distributed/autograd/init.cpp 7 8class DistAutogradContext: 9 def _context_id(self) -> int: ... 10 def _recv_functions(self) -> dict[int, Any]: ... 11 def _send_functions(self) -> dict[int, Any]: ... 12 def _known_worker_ids(self) -> set[int]: ... 13 14def _new_context() -> DistAutogradContext: ... 15def _release_context(context_id: int) -> None: ... 16def _get_max_id() -> int: ... 17def _is_valid_context(worker_id: int) -> bool: ... 18def _retrieve_context(context_id: int) -> DistAutogradContext: ... 19def _current_context() -> DistAutogradContext: ... 20def _init(worker_id: int) -> None: ... 21def _get_debug_info() -> dict[str, str]: ... 22def backward( 23 context_id: int, 24 roots: list[torch.Tensor], 25 retain_graph=False, 26) -> None: ... 27def get_gradients(context_id: int) -> dict[torch.Tensor, torch.Tensor]: ... 28