syntax = "proto3";

package tensorflow.data.experimental;

import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";

option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";

// 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.
message SnapshotRecord {
  repeated .tensorflow.TensorProto tensor = 1;
}

// This stores the metadata information present in each snapshot record.
message SnapshotMetadataRecord {
  // Stores the fingerprint of the graph that describes the dataset that is
  // snapshotted.
  string graph_hash = 1;
  // Run ID that this snapshot corresponds to.
  string run_id = 2;
  // Time when we started creating this snapshot.
  int64 creation_timestamp = 3;
  // Version of the snapshot data file format.
  int64 version = 4;
  // A list of tensor dtype corresponding to each element of the snapshot.
  repeated .tensorflow.DataType dtype = 5;
  // The number of elements in the snapshot.
  int64 num_elements = 6;

  bool finalized = 1000;
}

// Metadata for a single tensor in the Snapshot Record.
message TensorMetadata {
  .tensorflow.TensorShapeProto tensor_shape = 2;
  // Number of uncompressed bytes used to store the tensor representation.
  int64 tensor_size_bytes = 3;
}

// Metadata for all the tensors in a Snapshot Record.
message SnapshotTensorMetadata {
  repeated TensorMetadata tensor_metadata = 1;
}

// Metadata for a `tf.data.Dataset` distributed snapshot.
message DistributedSnapshotMetadata {
  // The element spec of the snapshotted dataset.
  bytes element_spec = 1;

  // Whether and how to compress the snapshot.  Supported values are defined in
  // `tsl::io::compression`.  In particular, an empty string specifies not to
  // compress.
  string compression = 2;
}
