syntax = "proto3";

option cc_enable_arenas = true;

package tensorflow;

// A TensorBundle addition which saves extra information about the objects which
// own variables, allowing for more robust checkpoint loading into modified
// programs.

message TrackableObjectGraph {
  message TrackableObject {
    message ObjectReference {
      // An index into `TrackableObjectGraph.nodes`, indicating the object
      // being referenced.
      int32 node_id = 1;
      // A user-provided name for the edge.
      string local_name = 2;
    }

    message SerializedTensor {
      // 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.
      string name = 1;
      // 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.
      string full_name = 2;
      // The generated name of the Tensor in the checkpoint.
      string checkpoint_key = 3;
      // Whether checkpoints should be considered as matching even without this
      // value restored. Used for non-critical values which don't affect the
      // TensorFlow graph, such as layer configurations.
      bool optional_restore = 4;
    }

    message SlotVariableReference {
      // An index into `TrackableObjectGraph.nodes`, indicating the
      // variable object this slot was created for.
      int32 original_variable_node_id = 1;
      // The name of the slot (e.g. "m"/"v").
      string slot_name = 2;
      // An index into `TrackableObjectGraph.nodes`, indicating the
      // `Object` with the value of the slot variable.
      int32 slot_variable_node_id = 3;
    }

    // Objects which this object depends on.
    repeated ObjectReference children = 1;
    // Serialized data specific to this object.
    repeated SerializedTensor attributes = 2;
    // Slot variables owned by this object.
    repeated SlotVariableReference slot_variables = 3;
  }

  repeated TrackableObject nodes = 1;
}
