basedpyright 1.36.2 → 1.37.0-961d59bacff1958014fe4e9ed2ef3116fe2266e2

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/core/protobuf/tpu/topology_pb2.pyi

@@ -0,0 +1,149 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+
+import builtins
+import collections.abc
+import sys
+import typing
+
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.internal.enum_type_wrapper
+import google.protobuf.message
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing.final
+class TPUHardwareFeature(google.protobuf.message.Message):
+    """Describes features of a tpu."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    class _EmbeddingFeature:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _EmbeddingFeatureEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[TPUHardwareFeature._EmbeddingFeature.ValueType], builtins.type
+    ):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        UNSUPPORTED: TPUHardwareFeature._EmbeddingFeature.ValueType  # 0
+        """No embedding lookup accelerator available on the tpu."""
+        V1: TPUHardwareFeature._EmbeddingFeature.ValueType  # 1
+        """Embedding lookup accelerator V1. The embedding lookup operation can only
+        be placed at the beginning of computation. Only one instance of embedding
+        lookup layer is allowed.
+        """
+        V2: TPUHardwareFeature._EmbeddingFeature.ValueType  # 2
+        """Embedding lookup accelerator V2. The embedding lookup operation can be
+        placed anywhere of the computation. Multiple instances of embedding
+        lookup layer is allowed.
+        """
+
+    class EmbeddingFeature(_EmbeddingFeature, metaclass=_EmbeddingFeatureEnumTypeWrapper):
+        """Embedding feature of a tpu."""
+
+    UNSUPPORTED: TPUHardwareFeature.EmbeddingFeature.ValueType  # 0
+    """No embedding lookup accelerator available on the tpu."""
+    V1: TPUHardwareFeature.EmbeddingFeature.ValueType  # 1
+    """Embedding lookup accelerator V1. The embedding lookup operation can only
+    be placed at the beginning of computation. Only one instance of embedding
+    lookup layer is allowed.
+    """
+    V2: TPUHardwareFeature.EmbeddingFeature.ValueType  # 2
+    """Embedding lookup accelerator V2. The embedding lookup operation can be
+    placed anywhere of the computation. Multiple instances of embedding
+    lookup layer is allowed.
+    """
+
+    EMBEDDING_FEATURE_FIELD_NUMBER: builtins.int
+    NUM_EMBEDDING_DEVICES_PER_CHIP_FIELD_NUMBER: builtins.int
+    embedding_feature: global___TPUHardwareFeature.EmbeddingFeature.ValueType
+    num_embedding_devices_per_chip: builtins.int
+    """Number of embedding accelerator devices per chip."""
+    def __init__(
+        self,
+        *,
+        embedding_feature: global___TPUHardwareFeature.EmbeddingFeature.ValueType | None = ...,
+        num_embedding_devices_per_chip: builtins.int | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "embedding_feature", b"embedding_feature", "num_embedding_devices_per_chip", b"num_embedding_devices_per_chip"
+        ],
+    ) -> None: ...
+
+global___TPUHardwareFeature = TPUHardwareFeature
+
+@typing.final
+class TopologyProto(google.protobuf.message.Message):
+    """Describes the geometry of a TPU mesh."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    MESH_SHAPE_FIELD_NUMBER: builtins.int
+    NUM_TASKS_FIELD_NUMBER: builtins.int
+    NUM_TPU_DEVICES_PER_TASK_FIELD_NUMBER: builtins.int
+    DEVICE_COORDINATES_FIELD_NUMBER: builtins.int
+    TPU_HARDWARE_FEATURE_FIELD_NUMBER: builtins.int
+    num_tasks: builtins.int
+    """Number of TensorFlow tasks in the cluster."""
+    num_tpu_devices_per_task: builtins.int
+    """Number of TPU devices per task."""
+    @property
+    def mesh_shape(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimensions of the TPU topology, in cores. Typically, this is a 4D
+        topology [x, y, z, core], where the major dimensions correspond to TPU
+        chips, and the minor dimension describes the number of cores on a multicore
+        chip.
+        """
+
+    @property
+    def device_coordinates(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """A flattened rank 3 int32 array with shape
+        [num_tasks, num_tpu_devices_per_task, len(mesh_shape)].
+        `tasks` is the number of tasks in the TPU cluster, `devices` is the number
+        of TPU devices per task, and the minor dimension corresponds to a position
+        in the TPU mesh topology. Each entry [task, device, axis] gives the
+        `axis`-th coordinate in the topology of a task/device pair.
+        """
+
+    @property
+    def tpu_hardware_feature(self) -> global___TPUHardwareFeature:
+        """TPU supported features."""
+
+    def __init__(
+        self,
+        *,
+        mesh_shape: collections.abc.Iterable[builtins.int] | None = ...,
+        num_tasks: builtins.int | None = ...,
+        num_tpu_devices_per_task: builtins.int | None = ...,
+        device_coordinates: collections.abc.Iterable[builtins.int] | None = ...,
+        tpu_hardware_feature: global___TPUHardwareFeature | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["tpu_hardware_feature", b"tpu_hardware_feature"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "device_coordinates",
+            b"device_coordinates",
+            "mesh_shape",
+            b"mesh_shape",
+            "num_tasks",
+            b"num_tasks",
+            "num_tpu_devices_per_task",
+            b"num_tpu_devices_per_task",
+            "tpu_hardware_feature",
+            b"tpu_hardware_feature",
+        ],
+    ) -> None: ...
+
+global___TopologyProto = TopologyProto

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/core/protobuf/tpu/tpu_embedding_configuration_pb2.pyi

@@ -0,0 +1,326 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+
+import builtins
+import collections.abc
+import sys
+import typing
+
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.internal.enum_type_wrapper
+import google.protobuf.message
+import tensorflow.core.protobuf.tpu.optimization_parameters_pb2
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing.final
+class TPUEmbeddingConfiguration(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    class _Mode:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _ModeEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[TPUEmbeddingConfiguration._Mode.ValueType], builtins.type
+    ):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        UNSPECIFIED: TPUEmbeddingConfiguration._Mode.ValueType  # 0
+        INFERENCE: TPUEmbeddingConfiguration._Mode.ValueType  # 1
+        TRAINING: TPUEmbeddingConfiguration._Mode.ValueType  # 2
+        BACKWARD_PASS_ONLY: TPUEmbeddingConfiguration._Mode.ValueType  # 3
+
+    class Mode(_Mode, metaclass=_ModeEnumTypeWrapper):
+        """Mode. Should the embedding layer program be run for inference (just forward
+        pass), training (both forward and backward pass) or just the backward_pass.
+        """
+
+    UNSPECIFIED: TPUEmbeddingConfiguration.Mode.ValueType  # 0
+    INFERENCE: TPUEmbeddingConfiguration.Mode.ValueType  # 1
+    TRAINING: TPUEmbeddingConfiguration.Mode.ValueType  # 2
+    BACKWARD_PASS_ONLY: TPUEmbeddingConfiguration.Mode.ValueType  # 3
+
+    class _ShardingStrategy:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _ShardingStrategyEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[TPUEmbeddingConfiguration._ShardingStrategy.ValueType],
+        builtins.type,
+    ):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        DIV_DEFAULT: TPUEmbeddingConfiguration._ShardingStrategy.ValueType  # 0
+        MOD: TPUEmbeddingConfiguration._ShardingStrategy.ValueType  # 1
+
+    class ShardingStrategy(_ShardingStrategy, metaclass=_ShardingStrategyEnumTypeWrapper):
+        """Sharding strategy of the embedding tables among the hosts.
+        If the sharding_strategy is "mod", each id is assigned to host
+        "id % num_hosts". For instance, 13 ids are split across 5 hosts as:
+        [[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]].
+        If the sharding_strategy is "div", ids are assigned to hosts in a
+        contiguous manner. In this case, 13 ids are split across 5 hosts as:
+        [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].
+        In both the strategies, if the id space does not evenly divide the number
+        of hosts, each of the first "table_descriptor.vocabulary_size % num_hosts"
+        hosts will be assigned one more id.
+        This partitioning strategy exactly follows that in the embedding_lookup
+        TensorFlow function at tensorflow/python/ops/embedding_ops.py.
+        """
+
+    DIV_DEFAULT: TPUEmbeddingConfiguration.ShardingStrategy.ValueType  # 0
+    MOD: TPUEmbeddingConfiguration.ShardingStrategy.ValueType  # 1
+
+    @typing.final
+    class TableDescriptor(google.protobuf.message.Message):
+        """Description of the various embedding tables."""
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        NAME_FIELD_NUMBER: builtins.int
+        VOCABULARY_SIZE_FIELD_NUMBER: builtins.int
+        DIMENSION_FIELD_NUMBER: builtins.int
+        NUM_FEATURES_FIELD_NUMBER: builtins.int
+        OPTIMIZATION_PARAMETERS_FIELD_NUMBER: builtins.int
+        name: builtins.str
+        """Name of the table."""
+        vocabulary_size: builtins.int
+        """Size of the vocabulary (i.e., number of rows) in the table."""
+        dimension: builtins.int
+        """The embedding dimension (i.e., the width of the embedding table)."""
+        num_features: builtins.int
+        """Number of features mapped to this table."""
+        @property
+        def optimization_parameters(self) -> tensorflow.core.protobuf.tpu.optimization_parameters_pb2.OptimizationParameters:
+            """Details of the learning algorithm used to update the embedding
+            parameters.
+            """
+
+        def __init__(
+            self,
+            *,
+            name: builtins.str | None = ...,
+            vocabulary_size: builtins.int | None = ...,
+            dimension: builtins.int | None = ...,
+            num_features: builtins.int | None = ...,
+            optimization_parameters: tensorflow.core.protobuf.tpu.optimization_parameters_pb2.OptimizationParameters | None = ...,
+        ) -> None: ...
+        def HasField(
+            self, field_name: typing.Literal["optimization_parameters", b"optimization_parameters"]
+        ) -> builtins.bool: ...
+        def ClearField(
+            self,
+            field_name: typing.Literal[
+                "dimension",
+                b"dimension",
+                "name",
+                b"name",
+                "num_features",
+                b"num_features",
+                "optimization_parameters",
+                b"optimization_parameters",
+                "vocabulary_size",
+                b"vocabulary_size",
+            ],
+        ) -> None: ...
+
+    @typing.final
+    class FeatureDescriptor(google.protobuf.message.Message):
+        """Description of different input features."""
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        NAME_FIELD_NUMBER: builtins.int
+        TABLE_ID_FIELD_NUMBER: builtins.int
+        INPUT_SHAPE_FIELD_NUMBER: builtins.int
+        name: builtins.str
+        """Name of the input feature."""
+        table_id: builtins.int
+        """Index of the corresponding table in the TableDescriptor list."""
+        @property
+        def input_shape(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+            """Static shape of the inputs (excluding the reduction axis). Note that
+            the shape of the actual inputs provided using the infeed op must be
+            strictly smaller than input_shape. The outputs received at the TensorCore
+            will have rank = input_shape.size() + 1. The innermost axis corresponds
+            to the embedding dimension. If the input has shape [m, n, k] (excluding
+            the reduction axis) and the embedding dimension is d, the output received
+            at the TensorCore will have shape [m, n, k, d].
+            """
+
+        def __init__(
+            self,
+            *,
+            name: builtins.str | None = ...,
+            table_id: builtins.int | None = ...,
+            input_shape: collections.abc.Iterable[builtins.int] | None = ...,
+        ) -> None: ...
+        def ClearField(
+            self, field_name: typing.Literal["input_shape", b"input_shape", "name", b"name", "table_id", b"table_id"]
+        ) -> None: ...
+
+    @typing.final
+    class SpmdSharding(google.protobuf.message.Message):
+        """SPMD (Single Program Multiple Data) sharding configuration for
+        TPUEmbedding. When model parallelism is used on the TensorCore, the number
+        of cores per replica must be passed to TPUEmbedding so that the right
+        shapes can be computed in the TF/XLA bridge.
+        """
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        ENABLED_FIELD_NUMBER: builtins.int
+        NUM_CORES_PER_REPLICA_FIELD_NUMBER: builtins.int
+        enabled: builtins.bool
+        """Whether SPMD sharding is enabled."""
+        num_cores_per_replica: builtins.int
+        """Number of cores per replica."""
+        def __init__(self, *, enabled: builtins.bool | None = ..., num_cores_per_replica: builtins.int | None = ...) -> None: ...
+        def ClearField(
+            self, field_name: typing.Literal["enabled", b"enabled", "num_cores_per_replica", b"num_cores_per_replica"]
+        ) -> None: ...
+
+    TABLE_DESCRIPTOR_FIELD_NUMBER: builtins.int
+    MODE_FIELD_NUMBER: builtins.int
+    BATCH_SIZE_PER_TENSOR_CORE_FIELD_NUMBER: builtins.int
+    NUM_HOSTS_FIELD_NUMBER: builtins.int
+    NUM_TENSOR_CORES_FIELD_NUMBER: builtins.int
+    SHARDING_STRATEGY_FIELD_NUMBER: builtins.int
+    PIPELINE_EXECUTION_WITH_TENSOR_CORE_FIELD_NUMBER: builtins.int
+    PROFILE_DATA_DIRECTORY_FIELD_NUMBER: builtins.int
+    FEATURE_DESCRIPTOR_FIELD_NUMBER: builtins.int
+    SPMD_SHARDING_FIELD_NUMBER: builtins.int
+    mode: global___TPUEmbeddingConfiguration.Mode.ValueType
+    batch_size_per_tensor_core: builtins.int
+    """Number of samples in each batch of embedding layer activations sent to
+    the TensorCore.
+    """
+    num_hosts: builtins.int
+    """Number of TPU hosts used for inference/training."""
+    num_tensor_cores: builtins.int
+    """Number of TensorCore used for inference/training."""
+    sharding_strategy: global___TPUEmbeddingConfiguration.ShardingStrategy.ValueType
+    pipeline_execution_with_tensor_core: builtins.bool
+    """This parameter determines if the execution of the sparse core will be
+    pipelined with that of the TensorCore. This parameter only affects results
+    when mode=TRAINING. If mode=INFERENCE or BACKWARD_PASS_ONLY, this parameter
+    does not affect execution and hence, is a don't care value.
+
+    false: The execution of the sparse core is not pipelined with that of the
+    TensorCore. The forward pass of every step on the sparse core is executed
+    only after the backward pass of the previous step is complete. And the
+    backward pass on the sparse core is executed only after the embedding
+    gradients have been computed on the TensorCore on every step. This ensures
+    that the activations on every step observe the gradient updates from the
+    previous step on both the sparse core and the TensorCore.
+
+    true: The execution of the sparse core is pipelined with that of the
+    TensorCore. The forward pass of every step on the sparse core can be
+    executed after the forward pass of the previous step is complete without
+    waiting for the backward pass. This improves the utilization of the sparse
+    core allowing it to process step N+1 while the embedding gradients for step
+    N are computed on the TensorCore. The backward pass of every step on the
+    sparse core is executed directly after the forward pass for the next step
+    is complete. The drawback is that embedding activations for step N+1 do not
+    observe the embedding gradient updates from step N. This could affect model
+    quality if step N and N+1 involve the same set of embedding IDs. However,
+    since the embedding updates are sparse, this is generally not considered a
+    problem.
+    """
+    profile_data_directory: builtins.str
+    """Directory where embedding lookup statistics are stored. These statistics
+    summarize information about the inputs to the embedding lookup
+    operation, in particular, the average number of embedding IDs per example
+    and how well the embedding IDs are load balanced across the system. The
+    lookup statistics are used during TPU initialization for embedding table
+    partitioning. Collection of lookup statistics is done at runtime by
+    profiling the embedding inputs: only 3% of input samples are profiled to
+    minimize host CPU overhead. Once a suitable number of samples are
+    profiled, the lookup statistics are saved to table-specific files in the
+    profile data directory generally at the end of a TPU training loop. The
+    filename corresponding to each table is obtained by hashing table specific
+    parameters (e.g., table name and number of features) and global
+    configuration parameters (e.g., sharding strategy and TPU worker task
+    count). The same profile data directory can be shared amongst several
+    models to reuse embedding lookup statistics.
+    """
+    @property
+    def table_descriptor(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+        global___TPUEmbeddingConfiguration.TableDescriptor
+    ]: ...
+    @property
+    def feature_descriptor(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+        global___TPUEmbeddingConfiguration.FeatureDescriptor
+    ]:
+        """If the feature_descriptor field is populated, the model should NOT populate
+        TableDescriptor.num_features and batch_size_per_tensor_core. These two
+        fields will be auto-populated by the TPUEmbedding rewrite passes.
+        """
+
+    @property
+    def spmd_sharding(self) -> global___TPUEmbeddingConfiguration.SpmdSharding: ...
+    def __init__(
+        self,
+        *,
+        table_descriptor: collections.abc.Iterable[global___TPUEmbeddingConfiguration.TableDescriptor] | None = ...,
+        mode: global___TPUEmbeddingConfiguration.Mode.ValueType | None = ...,
+        batch_size_per_tensor_core: builtins.int | None = ...,
+        num_hosts: builtins.int | None = ...,
+        num_tensor_cores: builtins.int | None = ...,
+        sharding_strategy: global___TPUEmbeddingConfiguration.ShardingStrategy.ValueType | None = ...,
+        pipeline_execution_with_tensor_core: builtins.bool | None = ...,
+        profile_data_directory: builtins.str | None = ...,
+        feature_descriptor: collections.abc.Iterable[global___TPUEmbeddingConfiguration.FeatureDescriptor] | None = ...,
+        spmd_sharding: global___TPUEmbeddingConfiguration.SpmdSharding | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["spmd_sharding", b"spmd_sharding"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "batch_size_per_tensor_core",
+            b"batch_size_per_tensor_core",
+            "feature_descriptor",
+            b"feature_descriptor",
+            "mode",
+            b"mode",
+            "num_hosts",
+            b"num_hosts",
+            "num_tensor_cores",
+            b"num_tensor_cores",
+            "pipeline_execution_with_tensor_core",
+            b"pipeline_execution_with_tensor_core",
+            "profile_data_directory",
+            b"profile_data_directory",
+            "sharding_strategy",
+            b"sharding_strategy",
+            "spmd_sharding",
+            b"spmd_sharding",
+            "table_descriptor",
+            b"table_descriptor",
+        ],
+    ) -> None: ...
+
+global___TPUEmbeddingConfiguration = TPUEmbeddingConfiguration
+
+@typing.final
+class TPUEmbeddingError(google.protobuf.message.Message):
+    """A placeholder message that is used to define a unique Status payload
+    URL for TPU embedding errors.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    def __init__(self) -> None: ...
+
+global___TPUEmbeddingError = TPUEmbeddingError

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/core/protobuf/trackable_object_graph_pb2.pyi

@@ -0,0 +1,213 @@
+"""
+@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
+import google.protobuf.wrappers_pb2
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing.final
+class TrackableObjectGraph(google.protobuf.message.Message):
+    """A TensorBundle addition which saves extra information about the objects which
+    own variables, allowing for more robust checkpoint loading into modified
+    programs.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.final
+    class TrackableObject(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        @typing.final
+        class ObjectReference(google.protobuf.message.Message):
+            DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+            NODE_ID_FIELD_NUMBER: builtins.int
+            LOCAL_NAME_FIELD_NUMBER: builtins.int
+            node_id: builtins.int
+            """An index into `TrackableObjectGraph.nodes`, indicating the object
+            being referenced.
+            """
+            local_name: builtins.str
+            """A user-provided name for the edge."""
+            def __init__(self, *, node_id: builtins.int | None = ..., local_name: builtins.str | None = ...) -> None: ...
+            def ClearField(self, field_name: typing.Literal["local_name", b"local_name", "node_id", b"node_id"]) -> None: ...
+
+        @typing.final
+        class SerializedTensor(google.protobuf.message.Message):
+            DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+            NAME_FIELD_NUMBER: builtins.int
+            FULL_NAME_FIELD_NUMBER: builtins.int
+            CHECKPOINT_KEY_FIELD_NUMBER: builtins.int
+            name: builtins.str
+            """A name for the Tensor. Simple variables have only one
+            `SerializedTensor` named "VARIABLE_VALUE" by convention. This value may
+            be restored on object creation as an optimization.
+            """
+            full_name: builtins.str
+            """The full name of the variable/tensor, if applicable. Used to allow
+            name-based loading of checkpoints which were saved using an
+            object-based API. Should match the checkpoint key which would have been
+            assigned by tf.train.Saver.
+            """
+            checkpoint_key: builtins.str
+            """The generated name of the Tensor in the checkpoint."""
+            def __init__(
+                self,
+                *,
+                name: builtins.str | None = ...,
+                full_name: builtins.str | None = ...,
+                checkpoint_key: builtins.str | None = ...,
+            ) -> None: ...
+            def ClearField(
+                self, field_name: typing.Literal["checkpoint_key", b"checkpoint_key", "full_name", b"full_name", "name", b"name"]
+            ) -> None: ...
+
+        @typing.final
+        class SlotVariableReference(google.protobuf.message.Message):
+            DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+            ORIGINAL_VARIABLE_NODE_ID_FIELD_NUMBER: builtins.int
+            SLOT_NAME_FIELD_NUMBER: builtins.int
+            SLOT_VARIABLE_NODE_ID_FIELD_NUMBER: builtins.int
+            original_variable_node_id: builtins.int
+            """An index into `TrackableObjectGraph.nodes`, indicating the
+            variable object this slot was created for.
+            """
+            slot_name: builtins.str
+            """The name of the slot (e.g. "m"/"v")."""
+            slot_variable_node_id: builtins.int
+            """An index into `TrackableObjectGraph.nodes`, indicating the
+            `Object` with the value of the slot variable.
+            """
+            def __init__(
+                self,
+                *,
+                original_variable_node_id: builtins.int | None = ...,
+                slot_name: builtins.str | None = ...,
+                slot_variable_node_id: builtins.int | None = ...,
+            ) -> None: ...
+            def ClearField(
+                self,
+                field_name: typing.Literal[
+                    "original_variable_node_id",
+                    b"original_variable_node_id",
+                    "slot_name",
+                    b"slot_name",
+                    "slot_variable_node_id",
+                    b"slot_variable_node_id",
+                ],
+            ) -> None: ...
+
+        CHILDREN_FIELD_NUMBER: builtins.int
+        ATTRIBUTES_FIELD_NUMBER: builtins.int
+        SLOT_VARIABLES_FIELD_NUMBER: builtins.int
+        REGISTERED_SAVER_FIELD_NUMBER: builtins.int
+        HAS_CHECKPOINT_VALUES_FIELD_NUMBER: builtins.int
+        @property
+        def children(
+            self,
+        ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+            global___TrackableObjectGraph.TrackableObject.ObjectReference
+        ]:
+            """Objects which this object depends on."""
+
+        @property
+        def attributes(
+            self,
+        ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+            global___TrackableObjectGraph.TrackableObject.SerializedTensor
+        ]:
+            """Serialized data specific to this object."""
+
+        @property
+        def slot_variables(
+            self,
+        ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+            global___TrackableObjectGraph.TrackableObject.SlotVariableReference
+        ]:
+            """Slot variables owned by this object."""
+
+        @property
+        def registered_saver(self) -> global___RegisteredSaver:
+            """The registered saver used to save this object. If this saver is not
+            present when loading the checkpoint, then loading will fail.
+            """
+
+        @property
+        def has_checkpoint_values(self) -> google.protobuf.wrappers_pb2.BoolValue:
+            """Whether this object has checkpoint values or descendants with checkpoint
+            values. This is computed at save time to avoid traversing the entire
+            object graph proto when restoring (which also has to traverse the live
+            object graph).
+            """
+
+        def __init__(
+            self,
+            *,
+            children: collections.abc.Iterable[global___TrackableObjectGraph.TrackableObject.ObjectReference] | None = ...,
+            attributes: collections.abc.Iterable[global___TrackableObjectGraph.TrackableObject.SerializedTensor] | None = ...,
+            slot_variables: (
+                collections.abc.Iterable[global___TrackableObjectGraph.TrackableObject.SlotVariableReference] | None
+            ) = ...,
+            registered_saver: global___RegisteredSaver | None = ...,
+            has_checkpoint_values: google.protobuf.wrappers_pb2.BoolValue | None = ...,
+        ) -> None: ...
+        def HasField(
+            self,
+            field_name: typing.Literal[
+                "has_checkpoint_values", b"has_checkpoint_values", "registered_saver", b"registered_saver"
+            ],
+        ) -> builtins.bool: ...
+        def ClearField(
+            self,
+            field_name: typing.Literal[
+                "attributes",
+                b"attributes",
+                "children",
+                b"children",
+                "has_checkpoint_values",
+                b"has_checkpoint_values",
+                "registered_saver",
+                b"registered_saver",
+                "slot_variables",
+                b"slot_variables",
+            ],
+        ) -> None: ...
+
+    NODES_FIELD_NUMBER: builtins.int
+    @property
+    def nodes(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___TrackableObjectGraph.TrackableObject]: ...
+    def __init__(
+        self, *, nodes: collections.abc.Iterable[global___TrackableObjectGraph.TrackableObject] | None = ...
+    ) -> None: ...
+    def ClearField(self, field_name: typing.Literal["nodes", b"nodes"]) -> None: ...
+
+global___TrackableObjectGraph = TrackableObjectGraph
+
+@typing.final
+class RegisteredSaver(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    NAME_FIELD_NUMBER: builtins.int
+    OBJECT_NAME_FIELD_NUMBER: builtins.int
+    name: builtins.str
+    """The name of the registered saver/restore function."""
+    object_name: builtins.str
+    """Unique auto-generated name of the object."""
+    def __init__(self, *, name: builtins.str | None = ..., object_name: builtins.str | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["name", b"name", "object_name", b"object_name"]) -> None: ...
+
+global___RegisteredSaver = RegisteredSaver