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";

// 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;
}