Add some functions from the config module.

Author: hoel-bagard

stubs/tensorflow/tensorflow/config/__init__.pyi

@@ -0,0 +1,11 @@
+from typing import NamedTuple
+
+from tensorflow.config import experimental as experimental
+
+class PhysicalDevice(NamedTuple):
+    name: str
+    device_type: str
+
+def list_physical_devices(device_type: None | str = None) -> list[PhysicalDevice]: ...
+def get_visible_devices(device_type: None | str = None) -> list[PhysicalDevice]: ...
+def set_visible_devices(devices: list[PhysicalDevice], device_type: None | str = None) -> None: ...

stubs/tensorflow/tensorflow/config/experimental.pyi

@@ -0,0 +1,13 @@
+from typing_extensions import TypedDict
+
+from tensorflow.config import PhysicalDevice
+
+class _MemoryInfo(TypedDict):
+    current: int
+    peak: int
+
+def get_memory_info(device: str) -> _MemoryInfo: ...
+def reset_memory_stats(device: str) -> None: ...
+def get_memory_usage(device: PhysicalDevice) -> bytes: ...
+def get_memory_growth(device: PhysicalDevice) -> bool: ...
+def set_memory_growth(device: PhysicalDevice, enable: bool) -> None: ...