/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
syntax = "proto3";

package tensorflow;

import "google/protobuf/wrappers.proto";

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

// 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;
      // Deprecated bool field for optional restore. This field has never been
      // set to True.
      reserved "optional_restore";
      reserved 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;

    // The registered saver used to save this object. If this saver is not
    // present when loading the checkpoint, then loading will fail.
    RegisteredSaver registered_saver = 4;

    // 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).
    google.protobuf.BoolValue has_checkpoint_values = 5;
  }

  repeated TrackableObject nodes = 1;
}

message RegisteredSaver {
  // The name of the registered saver/restore function.
  string name = 1;

  // Unique auto-generated name of the object.
  string object_name = 2;
}
