Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions stubs/tensorflow/tensorflow/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ class Operation:
def __getattr__(self, name: str) -> Incomplete: ...

class TensorShape(metaclass=ABCMeta):
__slots__ = ["_dims"]
def __init__(self, dims: ShapeLike) -> None: ...
@property
def rank(self) -> int: ...
Expand Down Expand Up @@ -308,6 +309,7 @@ class UnconnectedGradients(Enum):
_SpecProto = TypeVar("_SpecProto", bound=Message)

class TypeSpec(ABC, Generic[_SpecProto]):
__slots__ = ["_cached_cmp_key"]
@property
@abstractmethod
def value_type(self) -> Any: ...
Expand All @@ -323,6 +325,7 @@ class TypeSpec(ABC, Generic[_SpecProto]):
def most_specific_compatible_type(self, other: Self) -> Self: ...

class TensorSpec(TypeSpec[struct_pb2.TensorSpecProto]):
__slots__: list[str] = []
def __init__(self, shape: ShapeLike, dtype: DTypeLike = ..., name: str | None = None) -> None: ...
@property
def value_type(self) -> Tensor: ...
Expand All @@ -339,6 +342,7 @@ class TensorSpec(TypeSpec[struct_pb2.TensorSpecProto]):
def is_compatible_with(self, spec_or_tensor: Self | TensorCompatible) -> _bool: ... # type: ignore[override]

class SparseTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]):
__slots__ = ["_shape", "_dtype"]
def __init__(self, shape: ShapeLike | None = None, dtype: DTypeLike = ...) -> None: ...
@property
def value_type(self) -> SparseTensor: ...
Expand All @@ -350,6 +354,7 @@ class SparseTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]):
def from_value(cls, value: SparseTensor) -> Self: ...

class RaggedTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]):
__slots__ = ["_shape", "_dtype", "_ragged_rank", "_row_splits_dtype", "_flat_values_spec"]
def __init__(
self,
shape: ShapeLike | None = None,
Expand Down
1 change: 1 addition & 0 deletions stubs/tensorflow/tensorflow/dtypes.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ from tensorflow.python.framework.dtypes import HandleData
class _DTypeMeta(ABCMeta): ...

class DType(metaclass=_DTypeMeta):
__slots__ = ["_handle_data"]
def __init__(self, type_enum: int, handle_data: HandleData | None = None) -> None: ...
@property
def name(self) -> str: ...
Expand Down
7 changes: 7 additions & 0 deletions stubs/tensorflow/tensorflow/saved_model/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,13 @@ class Asset:
def __init__(self, path: str | Path | tf.Tensor) -> None: ...

class LoadOptions:
__slots__ = (
"allow_partial_checkpoint",
"experimental_io_device",
"experimental_skip_checkpoint",
"experimental_variable_policy",
"experimental_load_function_aliases",
)
allow_partial_checkpoint: bool
experimental_io_device: str | None
experimental_skip_checkpoint: bool
Expand Down
11 changes: 11 additions & 0 deletions stubs/tensorflow/tensorflow/train/__init__.pyi
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from _typeshed import Incomplete
from collections.abc import Callable
from typing import Any, TypeVar
from typing_extensions import Self
Expand All @@ -18,10 +19,20 @@ from tensorflow.python.trackable.base import Trackable
from tensorflow.python.training.tracking.autotrackable import AutoTrackable

class CheckpointOptions:
__slots__ = (
"experimental_io_device",
"experimental_enable_async_checkpoint",
"experimental_write_callbacks",
"enable_async",
"experimental_sharding_callback",
"experimental_skip_slot_variables",
)
experimental_io_device: None | str
experimental_enable_async_checkpoint: bool
experimental_write_callbacks: None | list[Callable[[str], object] | Callable[[], object]]
enable_async: bool
experimental_sharding_callback: Incomplete # should be ShardingCallback
experimental_skip_slot_variables: bool

def __init__(
self,
Expand Down