mirror of
https://github.com/davidhalter/typeshed.git
synced 2025-12-07 20:54:28 +08:00
tensorflow: Add some functions from the config module (#11325)
This commit is contained in:
13
stubs/tensorflow/tensorflow/config/__init__.pyi
Normal file
13
stubs/tensorflow/tensorflow/config/__init__.pyi
Normal 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: ...
|
||||
17
stubs/tensorflow/tensorflow/config/experimental.pyi
Normal file
17
stubs/tensorflow/tensorflow/config/experimental.pyi
Normal 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: ...
|
||||
Reference in New Issue
Block a user