"""
@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 tensorflow.core.framework.tensor_pb2
import tensorflow.core.framework.tensor_shape_pb2
import tensorflow.core.framework.types_pb2

DESCRIPTOR: google.protobuf.descriptor.FileDescriptor

@typing.final
class SnapshotRecord(google.protobuf.message.Message):
    """Each SnapshotRecord represents one batch of pre-processed input data. A batch
    consists of a list of tensors that we encode as TensorProtos. This message
    doesn't store the structure of the batch.
    """

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    TENSOR_FIELD_NUMBER: builtins.int
    @property
    def tensor(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[tensorflow.core.framework.tensor_pb2.TensorProto]: ...
    def __init__(
        self,
        *,
        tensor: collections.abc.Iterable[tensorflow.core.framework.tensor_pb2.TensorProto] | None = ...,
    ) -> None: ...
    def ClearField(self, field_name: typing.Literal["tensor", b"tensor"]) -> None: ...

global___SnapshotRecord = SnapshotRecord

@typing.final
class SnapshotMetadataRecord(google.protobuf.message.Message):
    """This stores the metadata information present in each snapshot record."""

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    GRAPH_HASH_FIELD_NUMBER: builtins.int
    RUN_ID_FIELD_NUMBER: builtins.int
    CREATION_TIMESTAMP_FIELD_NUMBER: builtins.int
    VERSION_FIELD_NUMBER: builtins.int
    DTYPE_FIELD_NUMBER: builtins.int
    NUM_ELEMENTS_FIELD_NUMBER: builtins.int
    FINALIZED_FIELD_NUMBER: builtins.int
    graph_hash: builtins.str
    """Stores the fingerprint of the graph that describes the dataset that is
    snapshotted.
    """
    run_id: builtins.str
    """Run ID that this snapshot corresponds to."""
    creation_timestamp: builtins.int
    """Time when we started creating this snapshot."""
    version: builtins.int
    """Version of the snapshot data file format."""
    num_elements: builtins.int
    """The number of elements in the snapshot."""
    finalized: builtins.bool
    @property
    def dtype(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[tensorflow.core.framework.types_pb2.DataType.ValueType]:
        """A list of tensor dtype corresponding to each element of the snapshot."""

    def __init__(
        self,
        *,
        graph_hash: builtins.str | None = ...,
        run_id: builtins.str | None = ...,
        creation_timestamp: builtins.int | None = ...,
        version: builtins.int | None = ...,
        dtype: collections.abc.Iterable[tensorflow.core.framework.types_pb2.DataType.ValueType] | None = ...,
        num_elements: builtins.int | None = ...,
        finalized: builtins.bool | None = ...,
    ) -> None: ...
    def ClearField(self, field_name: typing.Literal["creation_timestamp", b"creation_timestamp", "dtype", b"dtype", "finalized", b"finalized", "graph_hash", b"graph_hash", "num_elements", b"num_elements", "run_id", b"run_id", "version", b"version"]) -> None: ...

global___SnapshotMetadataRecord = SnapshotMetadataRecord

@typing.final
class TensorMetadata(google.protobuf.message.Message):
    """Metadata for a single tensor in the Snapshot Record."""

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    TENSOR_SHAPE_FIELD_NUMBER: builtins.int
    TENSOR_SIZE_BYTES_FIELD_NUMBER: builtins.int
    tensor_size_bytes: builtins.int
    """Number of uncompressed bytes used to store the tensor representation."""
    @property
    def tensor_shape(self) -> tensorflow.core.framework.tensor_shape_pb2.TensorShapeProto: ...
    def __init__(
        self,
        *,
        tensor_shape: tensorflow.core.framework.tensor_shape_pb2.TensorShapeProto | None = ...,
        tensor_size_bytes: builtins.int | None = ...,
    ) -> None: ...
    def HasField(self, field_name: typing.Literal["tensor_shape", b"tensor_shape"]) -> builtins.bool: ...
    def ClearField(self, field_name: typing.Literal["tensor_shape", b"tensor_shape", "tensor_size_bytes", b"tensor_size_bytes"]) -> None: ...

global___TensorMetadata = TensorMetadata

@typing.final
class SnapshotTensorMetadata(google.protobuf.message.Message):
    """Metadata for all the tensors in a Snapshot Record."""

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    TENSOR_METADATA_FIELD_NUMBER: builtins.int
    @property
    def tensor_metadata(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___TensorMetadata]: ...
    def __init__(
        self,
        *,
        tensor_metadata: collections.abc.Iterable[global___TensorMetadata] | None = ...,
    ) -> None: ...
    def ClearField(self, field_name: typing.Literal["tensor_metadata", b"tensor_metadata"]) -> None: ...

global___SnapshotTensorMetadata = SnapshotTensorMetadata

@typing.final
class DistributedSnapshotMetadata(google.protobuf.message.Message):
    """Metadata for a `tf.data.Dataset` distributed snapshot."""

    DESCRIPTOR: google.protobuf.descriptor.Descriptor

    ELEMENT_SPEC_FIELD_NUMBER: builtins.int
    COMPRESSION_FIELD_NUMBER: builtins.int
    element_spec: builtins.bytes
    """The element spec of the snapshotted dataset."""
    compression: builtins.str
    """Whether and how to compress the snapshot.  Supported values are defined in
    `tsl::io::compression`.  In particular, an empty string specifies not to
    compress.
    """
    def __init__(
        self,
        *,
        element_spec: builtins.bytes | None = ...,
        compression: builtins.str | None = ...,
    ) -> None: ...
    def ClearField(self, field_name: typing.Literal["compression", b"compression", "element_spec", b"element_spec"]) -> None: ...

global___DistributedSnapshotMetadata = DistributedSnapshotMetadata
