basedpyright 1.36.2 → 1.37.0-961d59bacff1958014fe4e9ed2ef3116fe2266e2

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/strings.pyi

@@ -0,0 +1,241 @@
+from collections.abc import Sequence
+from typing import Literal, TypeVar, overload
+
+from tensorflow import RaggedTensor, Tensor
+from tensorflow._aliases import StringTensorCompatible, TensorCompatible, UIntTensorCompatible
+from tensorflow.dtypes import DType
+
+_TensorOrRaggedTensor = TypeVar("_TensorOrRaggedTensor", Tensor, RaggedTensor)
+
+@overload
+def as_string(
+    input: TensorCompatible,
+    precision: int = -1,
+    scientific: bool = False,
+    shortest: bool = False,
+    width: int = -1,
+    fill: str = "",
+    name: str | None = None,
+) -> Tensor: ...
+@overload
+def as_string(
+    input: RaggedTensor,
+    precision: int = -1,
+    scientific: bool = False,
+    shortest: bool = False,
+    width: int = -1,
+    fill: str = "",
+    name: str | None = None,
+) -> RaggedTensor: ...
+def bytes_split(input: TensorCompatible | RaggedTensor, name: str | None = None) -> RaggedTensor: ...
+def format(
+    template: str, inputs: TensorCompatible, placeholder: str = "{}", summarize: int = 3, name: str | None = None
+) -> Tensor: ...
+def join(inputs: Sequence[TensorCompatible | RaggedTensor], separator: str = "", name: str | None = None) -> Tensor: ...
+@overload
+def length(input: TensorCompatible, unit: Literal["BYTE", "UTF8_CHAR"] = "BYTE", name: str | None = None) -> Tensor: ...
+@overload
+def length(input: RaggedTensor, unit: Literal["BYTE", "UTF8_CHAR"] = "BYTE", name: str | None = None) -> RaggedTensor: ...
+@overload
+def lower(input: TensorCompatible, encoding: Literal["utf-8", ""] = "", name: str | None = None) -> Tensor: ...
+@overload
+def lower(input: RaggedTensor, encoding: Literal["utf-8", ""] = "", name: str | None = None) -> RaggedTensor: ...
+def ngrams(
+    data: StringTensorCompatible | RaggedTensor,
+    ngram_width: int | Sequence[int],
+    separator: str = " ",
+    pad_values: tuple[int, int] | str | None = None,
+    padding_width: int | None = None,
+    preserve_short_sequences: bool = False,
+    name: str | None = None,
+) -> RaggedTensor: ...
+def reduce_join(
+    inputs: StringTensorCompatible | RaggedTensor,
+    axis: int | None = None,
+    keepdims: bool = False,
+    separator: str = "",
+    name: str | None = None,
+) -> Tensor: ...
+@overload
+def regex_full_match(input: StringTensorCompatible, pattern: StringTensorCompatible, name: str | None = None) -> Tensor: ...
+@overload
+def regex_full_match(input: RaggedTensor, pattern: StringTensorCompatible, name: str | None = None) -> RaggedTensor: ...
+@overload
+def regex_replace(
+    input: StringTensorCompatible,
+    pattern: StringTensorCompatible,
+    rewrite: StringTensorCompatible,
+    replace_global: bool = True,
+    name: str | None = None,
+) -> Tensor: ...
+@overload
+def regex_replace(
+    input: RaggedTensor,
+    pattern: StringTensorCompatible,
+    rewrite: StringTensorCompatible,
+    replace_global: bool = True,
+    name: str | None = None,
+) -> RaggedTensor: ...
+def split(
+    input: StringTensorCompatible | RaggedTensor,
+    sep: StringTensorCompatible | None = None,
+    maxsplit: int = -1,
+    name: str | None = None,
+) -> RaggedTensor: ...
+@overload
+def strip(input: StringTensorCompatible, name: str | None = None) -> Tensor: ...
+@overload
+def strip(input: RaggedTensor, name: str | None = None) -> RaggedTensor: ...
+@overload
+def substr(
+    input: StringTensorCompatible,
+    pos: TensorCompatible,
+    len: TensorCompatible,
+    unit: Literal["BYTE", "UTF8_CHAR"] = "BYTE",
+    name: str | None = None,
+) -> Tensor: ...
+@overload
+def substr(
+    input: RaggedTensor,
+    pos: TensorCompatible,
+    len: TensorCompatible,
+    unit: Literal["BYTE", "UTF8_CHAR"] = "BYTE",
+    name: str | None = None,
+) -> RaggedTensor: ...
+@overload
+def to_hash_bucket(input: StringTensorCompatible, num_buckets: int, name: str | None = None) -> Tensor: ...
+@overload
+def to_hash_bucket(input: RaggedTensor, num_buckets: int, name: str | None = None) -> RaggedTensor: ...
+@overload
+def to_hash_bucket_fast(input: StringTensorCompatible, num_buckets: int, name: str | None = None) -> Tensor: ...
+@overload
+def to_hash_bucket_fast(input: RaggedTensor, num_buckets: int, name: str | None = None) -> RaggedTensor: ...
+@overload
+def to_hash_bucket_strong(
+    input: StringTensorCompatible, num_buckets: int, key: Sequence[int], name: str | None = None
+) -> Tensor: ...
+@overload
+def to_hash_bucket_strong(input: RaggedTensor, num_buckets: int, key: Sequence[int], name: str | None = None) -> RaggedTensor: ...
+@overload
+def to_number(input: StringTensorCompatible, out_type: DType = ..., name: str | None = None) -> Tensor: ...
+@overload
+def to_number(input: RaggedTensor, out_type: DType = ..., name: str | None = None) -> RaggedTensor: ...
+@overload
+def unicode_decode(
+    input: StringTensorCompatible,
+    input_encoding: str,
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    replace_control_characters: bool = False,
+    name: str | None = None,
+) -> Tensor | RaggedTensor: ...
+@overload
+def unicode_decode(
+    input: RaggedTensor,
+    input_encoding: str,
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    replace_control_characters: bool = False,
+    name: str | None = None,
+) -> RaggedTensor: ...
+@overload
+def unicode_decode_with_offsets(
+    input: StringTensorCompatible,
+    input_encoding: str,
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    replace_control_characters: bool = False,
+    name: str | None = None,
+) -> tuple[_TensorOrRaggedTensor, _TensorOrRaggedTensor]: ...
+@overload
+def unicode_decode_with_offsets(
+    input: RaggedTensor,
+    input_encoding: str,
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    replace_control_characters: bool = False,
+    name: str | None = None,
+) -> tuple[RaggedTensor, RaggedTensor]: ...
+@overload
+def unicode_encode(
+    input: TensorCompatible,
+    output_encoding: Literal["UTF-8", "UTF-16-BE", "UTF-32-BE"],
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    name: str | None = None,
+) -> Tensor: ...
+@overload
+def unicode_encode(
+    input: RaggedTensor,
+    output_encoding: Literal["UTF-8", "UTF-16-BE", "UTF-32-BE"],
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    name: str | None = None,
+) -> RaggedTensor: ...
+@overload
+def unicode_script(input: TensorCompatible, name: str | None = None) -> Tensor: ...
+@overload
+def unicode_script(input: RaggedTensor, name: str | None = None) -> RaggedTensor: ...
+@overload
+def unicode_split(
+    input: StringTensorCompatible,
+    input_encoding: str,
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    name: str | None = None,
+) -> Tensor | RaggedTensor: ...
+@overload
+def unicode_split(
+    input: RaggedTensor,
+    input_encoding: str,
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    name: str | None = None,
+) -> RaggedTensor: ...
+@overload
+def unicode_split_with_offsets(
+    input: StringTensorCompatible,
+    input_encoding: str,
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    name: str | None = None,
+) -> tuple[_TensorOrRaggedTensor, _TensorOrRaggedTensor]: ...
+@overload
+def unicode_split_with_offsets(
+    input: RaggedTensor,
+    input_encoding: str,
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    name: str | None = None,
+) -> tuple[RaggedTensor, RaggedTensor]: ...
+@overload
+def unicode_transcode(
+    input: StringTensorCompatible,
+    input_encoding: str,
+    output_encoding: Literal["UTF-8", "UTF-16-BE", "UTF-32-BE"],
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    replace_control_characters: bool = False,
+    name: str | None = None,
+) -> Tensor: ...
+@overload
+def unicode_transcode(
+    input: RaggedTensor,
+    input_encoding: str,
+    output_encoding: Literal["UTF-8", "UTF-16-BE", "UTF-32-BE"],
+    errors: Literal["replace", "strict", "ignore"] = "replace",
+    replacement_char: int = 65533,
+    replace_control_characters: bool = False,
+    name: str | None = None,
+) -> RaggedTensor: ...
+def unsorted_segment_join(
+    inputs: StringTensorCompatible,
+    segment_ids: UIntTensorCompatible,
+    num_segments: UIntTensorCompatible,
+    separator: str = "",
+    name: str | None = None,
+) -> Tensor: ...
+@overload
+def upper(input: TensorCompatible, encoding: Literal["utf-8", ""] = "", name: str | None = None) -> Tensor: ...
+@overload
+def upper(input: RaggedTensor, encoding: Literal["utf-8", ""] = "", name: str | None = None) -> RaggedTensor: ...

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/summary.pyi

@@ -0,0 +1,58 @@
+import abc
+from collections.abc import Callable, Generator
+from contextlib import AbstractContextManager, contextmanager
+from typing import Literal
+from typing_extensions import Self
+
+import tensorflow as tf
+from tensorflow._aliases import FloatArray, IntArray
+from tensorflow.core.framework.graph_pb2 import GraphDef
+from tensorflow.experimental.dtensor import Mesh
+
+class SummaryWriter(metaclass=abc.ABCMeta):
+    def as_default(self, step: int | None = None) -> AbstractContextManager[Self]: ...
+    def close(self) -> None: ...
+    def flush(self) -> None: ...
+    def init(self) -> None: ...
+    def set_as_default(self, step: int | None = None) -> None: ...
+
+def audio(
+    name: str,
+    data: tf.Tensor,
+    sample_rate: int | tf.Tensor,
+    step: int | tf.Tensor | None = None,
+    max_outputs: int | tf.Tensor | None = 3,
+    encoding: Literal["wav"] | None = None,
+    description: str | None = None,
+) -> bool: ...
+def create_file_writer(
+    logdir: str,
+    max_queue: int | None = None,
+    flush_millis: int | None = None,
+    filename_suffix: str | None = None,
+    name: str | None = None,
+    experimental_trackable: bool = False,
+    experimental_mesh: Mesh | None = None,
+) -> SummaryWriter: ...
+def create_noop_writer() -> SummaryWriter: ...
+def flush(writer: SummaryWriter | None = None, name: str | None = None) -> tf.Operation: ...
+def graph(graph_data: tf.Graph | GraphDef) -> bool: ...
+def histogram(
+    name: str, data: tf.Tensor, step: int | None = None, buckets: int | None = None, description: str | None = None
+) -> bool: ...
+def image(
+    name: str,
+    data: tf.Tensor | FloatArray | IntArray,
+    step: int | tf.Tensor | None = None,
+    max_outputs: int | None = 3,
+    description: str | None = None,
+) -> bool: ...
+@contextmanager
+def record_if(condition: bool | tf.Tensor | Callable[[], bool]) -> Generator[None]: ...
+def scalar(name: str, data: float | tf.Tensor, step: int | tf.Tensor | None = None, description: str | None = None) -> bool: ...
+def should_record_summaries() -> bool: ...
+def text(name: str, data: str | tf.Tensor, step: int | tf.Tensor | None = None, description: str | None = None) -> bool: ...
+def trace_export(name: str, step: int | tf.Tensor | None = None, profiler_outdir: str | None = None) -> None: ...
+def trace_off() -> None: ...
+def trace_on(graph: bool = True, profiler: bool = False, profiler_outdir: str | None = None) -> None: ...
+def write(tag: str, tensor: tf.Tensor, step: int | tf.Tensor | None = None, metadata=None, name: str | None = None) -> bool: ...

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/train/__init__.pyi

@@ -0,0 +1,74 @@
+from collections.abc import Callable
+from typing import Any, TypeVar
+from typing_extensions import Self
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.core.example.example_pb2 import Example as Example
+from tensorflow.core.example.feature_pb2 import (
+    BytesList as BytesList,
+    Feature as Feature,
+    Features as Features,
+    FloatList as FloatList,
+    Int64List as Int64List,
+)
+from tensorflow.core.protobuf.cluster_pb2 import ClusterDef as ClusterDef
+from tensorflow.core.protobuf.tensorflow_server_pb2 import ServerDef as ServerDef
+from tensorflow.python.trackable.base import Trackable
+from tensorflow.python.training.tracking.autotrackable import AutoTrackable
+
+class CheckpointOptions:
+    experimental_io_device: None | str
+    experimental_enable_async_checkpoint: bool
+    experimental_write_callbacks: None | list[Callable[[str], object] | Callable[[], object]]
+    enable_async: bool
+
+    def __init__(
+        self,
+        experimental_io_device: None | str = None,
+        experimental_enable_async_checkpoint: bool = False,
+        experimental_write_callbacks: None | list[Callable[[str], object] | Callable[[], object]] = None,
+        enable_async: bool = False,
+        experimental_skip_slot_variables: bool = False,
+        experimental_sharding_callback=None,
+    ) -> None: ...
+
+_T = TypeVar("_T", bound=list[str] | tuple[str] | dict[int, str])
+
+class ClusterSpec:
+    def __init__(self, cluster: dict[str, _T] | ClusterDef | ClusterSpec) -> None: ...
+    def as_dict(self) -> dict[str, list[str] | tuple[str] | dict[int, str]]: ...
+    def num_tasks(self, job_name: str) -> int: ...
+
+class _CheckpointLoadStatus:
+    def assert_consumed(self) -> Self: ...
+    def assert_existing_objects_matched(self) -> Self: ...
+    def assert_nontrivial_match(self) -> Self: ...
+    def expect_partial(self) -> Self: ...
+
+class Checkpoint(AutoTrackable):
+    def __init__(self, root: Trackable | None = None, **kwargs: Trackable) -> None: ...
+    def read(self, save_path: str, options: CheckpointOptions | None = None) -> _CheckpointLoadStatus: ...
+    def restore(self, save_path: str, options: CheckpointOptions | None = None) -> _CheckpointLoadStatus: ...
+    def save(self, file_prefix: str, options: CheckpointOptions | None = None) -> str: ...
+    def sync(self) -> None: ...
+    def write(self, file_prefix: str, options: CheckpointOptions | None = None) -> str: ...
+
+class CheckpointManager:
+    def __init__(
+        self,
+        checkpoint: Checkpoint,
+        directory: str,
+        max_to_keep: int,
+        keep_checkpoint_every_n_hours: int | None = None,
+        checkpoint_name: str = "ckpt",
+        step_counter: tf.Variable | None = None,
+        checkpoint_interval: int | None = None,
+        init_fn: Callable[[], object] | None = None,
+    ) -> None: ...
+    def _sweep(self) -> None: ...
+
+def latest_checkpoint(checkpoint_dir: str, latest_filename: str | None = None) -> str: ...
+def load_variable(ckpt_dir_or_file: str, name: str) -> np.ndarray[Any, Any]: ...
+def list_variables(ckpt_dir_or_file: str) -> list[tuple[str, list[int]]]: ...
+def __getattr__(name: str): ...  # incomplete module

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/train/experimental.pyi

@@ -0,0 +1,12 @@
+import abc
+from typing_extensions import Self
+
+from tensorflow.python.trackable.base import Trackable
+
+class PythonState(Trackable, metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def serialize(self) -> str: ...
+    @abc.abstractmethod
+    def deserialize(self, string_value: str) -> Self: ...
+
+def __getattr__(name: str): ...  # incomplete module

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/tsl/protobuf/coordination_config_pb2.pyi

@@ -0,0 +1,156 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+
+import builtins
+import collections.abc
+import typing
+
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.message
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing.final
+class CoordinatedJob(google.protobuf.message.Message):
+    """Represents a job type and the number of tasks under this job.
+    For example, ("worker", 20) implies that there will be 20 worker tasks.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    NAME_FIELD_NUMBER: builtins.int
+    NUM_TASKS_FIELD_NUMBER: builtins.int
+    name: builtins.str
+    num_tasks: builtins.int
+    def __init__(self, *, name: builtins.str | None = ..., num_tasks: builtins.int | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["name", b"name", "num_tasks", b"num_tasks"]) -> None: ...
+
+global___CoordinatedJob = CoordinatedJob
+
+@typing.final
+class CoordinationServiceConfig(google.protobuf.message.Message):
+    """Coordination service configuration parameters.
+    The system picks appropriate values for fields that are not set.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    SERVICE_TYPE_FIELD_NUMBER: builtins.int
+    SERVICE_LEADER_FIELD_NUMBER: builtins.int
+    ENABLE_HEALTH_CHECK_FIELD_NUMBER: builtins.int
+    CLUSTER_REGISTER_TIMEOUT_IN_MS_FIELD_NUMBER: builtins.int
+    HEARTBEAT_TIMEOUT_IN_MS_FIELD_NUMBER: builtins.int
+    COORDINATED_JOB_LIST_FIELD_NUMBER: builtins.int
+    SHUTDOWN_BARRIER_TIMEOUT_IN_MS_FIELD_NUMBER: builtins.int
+    AGENT_DESTRUCTION_WITHOUT_SHUTDOWN_FIELD_NUMBER: builtins.int
+    RECOVERABLE_JOBS_FIELD_NUMBER: builtins.int
+    ALLOW_NEW_INCARNATION_TO_RECONNECT_FIELD_NUMBER: builtins.int
+    FORCE_DISABLE_FIELD_NUMBER: builtins.int
+    POLL_FOR_ERROR_FROM_SERVICE_AT_STARTUP_FIELD_NUMBER: builtins.int
+    service_type: builtins.str
+    """Type of coordination service implementation to enable.
+    For example, setting the service type as "standalone" starts a service
+    instance on the leader task to provide the coordination services such as
+    heartbeats and consistent key-value store.
+    """
+    service_leader: builtins.str
+    """Address where the coordination service instance is hosted."""
+    enable_health_check: builtins.bool
+    """Whether to enable the health check mechanism."""
+    cluster_register_timeout_in_ms: builtins.int
+    """Maximum wait time for all members in the cluster to be registered."""
+    heartbeat_timeout_in_ms: builtins.int
+    """Heartbeat timeout, if a task does not record heartbeat in this time
+    window, it will be considered disconnected.
+    Note: This is also used as a grace period to accept any heartbeats after
+    the agent has disconnected, to account for the lag time between the service
+    recording the state change and the agent stopping heartbeats.
+    """
+    shutdown_barrier_timeout_in_ms: builtins.int
+    """Denotes how long to wait for all coordination agents to reach the barriers
+    (after the first shutdown request) before disconnecting together. If
+    set to 0, no barrier is imposed upon shutdown and each worker can
+    disconnect individually.
+    """
+    agent_destruction_without_shutdown: builtins.bool
+    """If set, agents do not make an explicit Shutdown() call. Service will only
+    find out about the disconnecte agent via stale heartbeats. Used for
+    testing.
+    """
+    allow_new_incarnation_to_reconnect: builtins.bool
+    """If a task restarts with a new incarnation, we may allow it to reconnect
+    silently. This is useful when we know that a task can immediately resume
+    work upon re-connecting to the service.
+    """
+    force_disable: builtins.bool
+    """Disables coordination service.
+    Some libraries enable coordination service by default even if the user did
+    not specify any config. This field allows users to explicitly disable
+    coordination service under all situations.
+    """
+    poll_for_error_from_service_at_startup: builtins.bool
+    """Use long polling to get error from coordination service as the error
+    propagation mechanism.
+    """
+    @property
+    def coordinated_job_list(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___CoordinatedJob]: ...
+    @property
+    def recoverable_jobs(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]:
+        """The list of jobs which are recoverable. If a task in this list fails,
+        it will not propagate error to other tasks.
+        If empty, no jobs will be recoverable and every task failure will cause
+        error propagation to other tasks.
+        """
+
+    def __init__(
+        self,
+        *,
+        service_type: builtins.str | None = ...,
+        service_leader: builtins.str | None = ...,
+        enable_health_check: builtins.bool | None = ...,
+        cluster_register_timeout_in_ms: builtins.int | None = ...,
+        heartbeat_timeout_in_ms: builtins.int | None = ...,
+        coordinated_job_list: collections.abc.Iterable[global___CoordinatedJob] | None = ...,
+        shutdown_barrier_timeout_in_ms: builtins.int | None = ...,
+        agent_destruction_without_shutdown: builtins.bool | None = ...,
+        recoverable_jobs: collections.abc.Iterable[builtins.str] | None = ...,
+        allow_new_incarnation_to_reconnect: builtins.bool | None = ...,
+        force_disable: builtins.bool | None = ...,
+        poll_for_error_from_service_at_startup: builtins.bool | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "agent_destruction_without_shutdown",
+            b"agent_destruction_without_shutdown",
+            "allow_new_incarnation_to_reconnect",
+            b"allow_new_incarnation_to_reconnect",
+            "cluster_register_timeout_in_ms",
+            b"cluster_register_timeout_in_ms",
+            "coordinated_job_list",
+            b"coordinated_job_list",
+            "enable_health_check",
+            b"enable_health_check",
+            "force_disable",
+            b"force_disable",
+            "heartbeat_timeout_in_ms",
+            b"heartbeat_timeout_in_ms",
+            "poll_for_error_from_service_at_startup",
+            b"poll_for_error_from_service_at_startup",
+            "recoverable_jobs",
+            b"recoverable_jobs",
+            "service_leader",
+            b"service_leader",
+            "service_type",
+            b"service_type",
+            "shutdown_barrier_timeout_in_ms",
+            b"shutdown_barrier_timeout_in_ms",
+        ],
+    ) -> None: ...
+
+global___CoordinationServiceConfig = CoordinationServiceConfig