/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed 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.*/
// Modification Copyright (C) 2024 Intel Corporation

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