tensorflow: Add some functions from the config module (#11325)

This commit is contained in:
Hoël Bagard
2024-02-01 00:59:45 +09:00
committed by GitHub
parent faf9d77d79
commit 284afde9cc
2 changed files with 30 additions and 0 deletions

View File

@@ -0,0 +1,13 @@
from _typeshed import Incomplete
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] | PhysicalDevice, device_type: None | str = None) -> None: ...
def __getattr__(name: str) -> Incomplete: ...

View File

@@ -0,0 +1,17 @@
import typing_extensions
from _typeshed import Incomplete
from typing 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: ...
@typing_extensions.deprecated("This function is deprecated in favor of tf.config.experimental.get_memory_info")
def get_memory_usage(device: PhysicalDevice) -> int: ...
def get_memory_growth(device: PhysicalDevice) -> bool: ...
def set_memory_growth(device: PhysicalDevice, enable: bool) -> None: ...
def __getattr__(name: str) -> Incomplete: ...