/* 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/any.proto";
import "ov_tensorflow/tensor_shape.proto";
import "ov_tensorflow/types.proto";
import "ov_tensorflow/variable.proto";
import "ov_tensorflow/versions.proto";
import "ov_tensorflow/struct.proto";
import "ov_tensorflow/trackable_object_graph.proto";

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

// A SavedObjectGraph is part of object-based SavedModels in TF 2.0. It
// describes the directed graph of Python objects (or equivalent in other
// languages) that make up a model, with nodes[0] at the root.

// SavedObjectGraph shares some structure with TrackableObjectGraph, but
// SavedObjectGraph belongs to the MetaGraph and contains pointers to functions
// and type information, while TrackableObjectGraph lives in the checkpoint
// and contains pointers only to variable values.

message SavedObjectGraph {
  // Flattened list of objects in the object graph.
  //
  // The position of the object in this list indicates its id.
  // Nodes[0] is considered the root node.
  repeated SavedObject nodes = 1;

  // Information about captures and output structures in concrete functions.
  // Referenced from SavedBareConcreteFunction and SavedFunction.
  map<string, SavedConcreteFunction> concrete_functions = 2;
}

message SavedObject {
  // Objects which this object depends on: named edges in the dependency
  // graph.
  //
  // Note: All kinds of SavedObject may have children, except
  // "constant" and "captured_tensor".
  repeated TrackableObjectGraph.TrackableObject.ObjectReference children = 1;

  // Ordered list of dependencies that must be loaded before this object.
  // SavedModel loads with the bottom-up approach, by first creating all objects
  // (in the order defined by the dependencies), then connecting the edges.
  repeated TrackableObjectGraph.TrackableObject.ObjectReference dependencies =
      15;

  // Removed when forking SavedObject from TrackableObjectGraph.
  reserved "attributes";
  reserved 2;

  // Slot variables owned by this object. This describes the three-way
  // (optimizer, variable, slot variable) relationship; none of the three
  // depend on the others directly.
  //
  // Note: currently only valid if kind == "user_object".
  repeated TrackableObjectGraph.TrackableObject.SlotVariableReference
      slot_variables = 3;

  oneof kind {
    SavedUserObject user_object = 4;
    SavedAsset asset = 5;
    SavedFunction function = 6;
    SavedVariable variable = 7;
    SavedBareConcreteFunction bare_concrete_function = 8;
    SavedConstant constant = 9;
    SavedResource resource = 10;
    CapturedTensor captured_tensor = 12;
  }

  // Stores the functions used to save and restore this object. At most one of
  // `saveable_objects` or `registered_saver` is defined for each SavedObject.
  // See the comment below for the difference between SaveableObject and
  // registered savers.
  map<string, SaveableObject> saveable_objects = 11;

  // The fields below are filled when the user serializes a registered Trackable
  // class or an object with a registered saver function.
  //
  // Registered classes may save additional metadata and supersede the
  // default loading process where nodes are recreated from the proto.
  // If the registered class cannot be found, then the object will load as one
  // one of the default trackable objects: Autotrackable (a class similar to
  // tf.Module), tf.function, or tf.Variable.
  //
  // Unlike SaveableObjects, which store the functions for saving and restoring
  // from tensors, registered savers allow Trackables to write checkpoint shards
  // directly (e.g. for performance or coordination reasons).
  // *All registered savers must be available when loading the SavedModel.*

  // The name of the registered class of the form "{package}.{class_name}".
  // This field is used to search for the registered class at loading time.
  string registered_name = 13;
  // The user-generated proto storing metadata for this object, to be passed to
  // the registered classes's _deserialize_from_proto method when this object is
  // loaded from the SavedModel.
  google.protobuf.Any serialized_user_proto = 14;

  // String name of the registered saver. At most one of `saveable_objects` or
  // `registered_saver` is defined for each SavedObject.
  string registered_saver = 16;
}

// A SavedUserObject is an object (in the object-oriented language of the
// TensorFlow program) of some user- or framework-defined class other than
// those handled specifically by the other kinds of SavedObjects.
//
// This object cannot be evaluated as a tensor, and therefore cannot be bound
// to an input of a function.
message SavedUserObject {
  // Corresponds to a registration of the type to use in the loading program.
  string identifier = 1;
  // Version information from the producer of this SavedUserObject.
  VersionDef version = 2;
  // Metadata for deserializing this object.
  //
  // Deprecated! At the time of deprecation, Keras was the only user of this
  // field, and its saving and loading code will be updated shortly.
  // Please save your application-specific metadata to a separate file.
  string metadata = 3 [deprecated = true];
}

// A SavedAsset points to an asset in the MetaGraph.
//
// When bound to a function this object evaluates to a tensor with the absolute
// filename. Users should not depend on a particular part of the filename to
// remain stable (e.g. basename could be changed).
message SavedAsset {
  // Index into `MetaGraphDef.asset_file_def[]` that describes the Asset.
  //
  // Only the field `AssetFileDef.filename` is used. Other fields, such as
  // `AssetFileDef.tensor_info`, MUST be ignored.
  int32 asset_file_def_index = 1;
}

// A function with multiple signatures, possibly with non-Tensor arguments.
message SavedFunction {
  repeated string concrete_functions = 1;
  FunctionSpec function_spec = 2;
}

message CapturedTensor {
  // Name of captured tensor
  string name = 1;

  // Name of concrete function which contains the computed graph tensor.
  string concrete_function = 2;
}

// Stores low-level information about a concrete function. Referenced in either
// a SavedFunction or a SavedBareConcreteFunction.
message SavedConcreteFunction {
  repeated int32 bound_inputs = 2;

  // Input in canonicalized form that was received to create this concrete
  // function.
  StructuredValue canonicalized_input_signature = 3;
  // Output that was the return value of this function after replacing all
  // Tensors with TensorSpecs. This can be an arbitrary nested function and will
  // be used to reconstruct the full structure from pure tensors.
  StructuredValue output_signature = 4;
}

message SavedBareConcreteFunction {
  // Identifies a SavedConcreteFunction.
  string concrete_function_name = 1;

  // A sequence of unique strings, one per Tensor argument.
  repeated string argument_keywords = 2;
  // The prefix of `argument_keywords` which may be identified by position.
  int64 allowed_positional_arguments = 3;
  // The spec of the function that this ConcreteFunction is traced from. This
  // allows the ConcreteFunction to be called with nest structure inputs. This
  // field may not be populated. If this field is absent, the concrete function
  // can only be called with flat inputs.
  // TODO(b/169361281): support calling saved ConcreteFunction with structured
  // inputs in C++ SavedModel API.
  FunctionSpec function_spec = 4;
}

message SavedConstant {
  // An Operation name for a ConstantOp in this SavedObjectGraph's MetaGraph.
  string operation = 1;
}

// Represents a Variable that is initialized by loading the contents from the
// checkpoint.
message SavedVariable {
  DataType dtype = 1;
  TensorShapeProto shape = 2;
  bool trainable = 3;
  VariableSynchronization synchronization = 4;
  VariableAggregation aggregation = 5;
  string name = 6;
  string device = 7;
  // List of component variables for a distributed variable.
  //
  // When this field is non-empty, the SavedVariable will be assumed
  // to be a distributed variable defined by the components listed here.
  //
  // This is only supported by experimental loaders at the moment.
  repeated SavedVariable experimental_distributed_variable_components = 8;
}

// Represents `FunctionSpec` used in `Function`. This represents a
// function that has been wrapped as a TensorFlow `Function`.
message FunctionSpec {
  // Full arg spec from inspect.getfullargspec().
  StructuredValue fullargspec = 1;
  // Whether this represents a class method.
  bool is_method = 2;
  // The input signature, if specified.
  StructuredValue input_signature = 5;

  // Whether the function should be compiled by XLA.
  //
  // The public interface to `tf.function` uses an optional boolean to
  // represent three distinct states for this field.  Unfortunately, proto3
  // removes the ability to explicitly check for the presence or absence of a
  // field, so we instead map to an enum.
  //
  // See `tf.function` for details.
  enum JitCompile {
    DEFAULT = 0;
    ON = 1;
    OFF = 2;
  }
  JitCompile jit_compile = 6;

  reserved 3, 4;
}

// A SavedResource represents a TF object that holds state during its lifetime.
// An object of this type can have a reference to a:
// create_resource() and an initialize() function.
message SavedResource {
  // A device specification indicating a required placement for the resource
  // creation function, e.g. "CPU". An empty string allows the user to select a
  // device.
  string device = 1;
}

message SaveableObject {
  // Node ids of concrete functions for saving and loading from a checkpoint.
  // These functions save and restore directly from tensors.
  int32 save_function = 2;
  int32 restore_function = 3;
}
