[BE]: Backport runtime_checkable perf improvements/behavior from 3.12 (#155130)

Backports some behavior changes and performance improvements with runtime_checkable in 3.12 to older versions of Python. Should be free performance improvement on typing checking protocols since everything works on Python 3.12.

The difference between the two versions of runtime_checkable is [these lines](40e22ebb2c/src/typing_extensions.py (L800-L823)).

Pull Request resolved: https://github.com/pytorch/pytorch/pull/155130
Approved by: https://github.com/rec, https://github.com/aorenste

torch/_C/__init__.pyi.in

@@ -16,12 +16,10 @@ from typing import (
     Literal,
     NamedTuple,
     overload,
-    Protocol,
-    runtime_checkable,
     SupportsIndex,
     TypeVar,
 )
-from typing_extensions import ParamSpec, Self, TypeAlias
+from typing_extensions import ParamSpec, Protocol, runtime_checkable, Self, TypeAlias
 
 import numpy
 

torch/distributed/_checkpointable.py

@@ -1,5 +1,5 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
-from typing import Protocol, runtime_checkable
+from typing_extensions import Protocol, runtime_checkable
 
 import torch
 

torch/distributed/checkpoint/staging.py

@@ -1,5 +1,5 @@
-from typing import Optional, runtime_checkable
-from typing_extensions import Protocol
+from typing import Optional
+from typing_extensions import Protocol, runtime_checkable
 
 from torch.distributed._state_dict_utils import _copy_state_dict, _create_cpu_state_dict
 from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE

torch/distributed/checkpoint/stateful.py

@@ -1,5 +1,5 @@
-from typing import Any, runtime_checkable, TypeVar
-from typing_extensions import Protocol
+from typing import Any, TypeVar
+from typing_extensions import Protocol, runtime_checkable
 
 
 __all__ = ["Stateful", "StatefulT"]

torch/onnx/_internal/fx/type_utils.py

@@ -4,7 +4,8 @@
 from __future__ import annotations
 
 from collections.abc import Mapping, Sequence
-from typing import Any, Optional, Protocol, runtime_checkable, TYPE_CHECKING, Union
+from typing import Any, Optional, TYPE_CHECKING, Union
+from typing_extensions import Protocol, runtime_checkable
 
 import onnx
 

torch/onnx/_internal/io_adapter.py

@@ -1,7 +1,8 @@
 # mypy: allow-untyped-defs
 from __future__ import annotations
 
-from typing import Any, Callable, Protocol, runtime_checkable, TYPE_CHECKING
+from typing import Any, Callable, TYPE_CHECKING
+from typing_extensions import Protocol, runtime_checkable
 
 import torch
 import torch.export as torch_export

torch/utils/benchmark/utils/_stubs.py

@@ -1,4 +1,5 @@
-from typing import Any, Callable, Protocol, runtime_checkable
+from typing import Any, Callable
+from typing_extensions import Protocol, runtime_checkable
 
 
 class TimerClass(Protocol):