syntax = "proto3";

package tensorflow;
option cc_enable_arenas = true;
option java_outer_classname = "FunctionProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.framework";

import "tensorflow/core/framework/attr_value.proto";
import "tensorflow/core/framework/node_def.proto";
import "tensorflow/core/framework/op_def.proto";

// A library is a set of named functions.
message FunctionDefLibrary {
  repeated FunctionDef function = 1;
  repeated GradientDef gradient = 2;
}

// A function can be instantiated when the runtime can bind every attr
// with a value. When a GraphDef has a call to a function, it must
// have binding for every attr defined in the signature.
//
// TODO(zhifengc):
//   * device spec, etc.
message FunctionDef {
  // The definition of the function's name, arguments, return values,
  // attrs etc.
  OpDef signature = 1;

  // TO BE REPLACED

  // The body of the function.
  repeated Node node = 2;  // function.node.ret[*] are unique.

  // A node is a multi-value assignment:
  //   (ret[0], ret[1], ...) = func(arg[0], arg[1], ...)
  //
  // By convention, "func" is resolved by consulting with a user-defined
  // library first. If not resolved, "func" is assumed to be a builtin op.
  message Node {
    // This node produces multiple outputs. They are named ret[0],
    // ret[1], ..., etc.
    //
    // REQUIRES: function.node.ret[*] are unique across all nodes.
    // REQUIRES: ret.size == func/op def's number of output args.
    repeated string ret = 1;

    // The op/function name.
    string op = 2;

    // Arguments passed to this func/op.
    //
    // arg[i] must be either one of
    // function.signature.input_args[*].name or one of
    // function.node[*].ret[*].
    //
    // REQUIRES: arg.size == func/op def's number of input args.
    repeated string arg = 3;

    // Control dependencies.
    //
    // dep[i] must be one of function.node[*].ret[*] or one of
    // function.signature.input_args[*].name.
    repeated string dep = 4;

    // Attrs.
    //
    // 'attr' maps names defined by 'func's attr defs to attr values.
    // attr values may have placeholders which are substituted
    // recursively by concrete values when this node is instantiated.
    // These placeholders must name an attr listed in the FunctionDef's
    // signature.
    map<string, AttrValue> attr = 5;
  }

  // WILL REPLACE THE ABOVE

  // If node_def is present, and the consumer is at GraphDef version
  // >= 12, then these fields are used and `node` is ignored.  If the
  // consumer's GraphDef version is < 12 or this field is empty, then
  // `node` is used.  This allows producers to fill both fields to
  // remain compatible with old consumers.  At some future GraphDef
  // version, `node` will be ignored even if `node_def` is empty.
  // TODO(josh11b): Finish this transition.

  // In both of the following fields, there is the need to specify an
  // output that is used as either the input to another node (in
  // `node_def`) or as a return value of the function (in `ret`).
  // Unlike the NodeDefs in GraphDef, we need to be able to specify a
  // list in some cases (instead of just single outputs).  Also, we
  // need to be able to deal with lists of unknown length (so the
  // output index may not be known at function definition time).  So
  // we use the following format instead:
  // * "fun_in" where "fun_in" is the name of a function input arg in
  //   the `signature` field above.  This represents that input, whether
  //   it is a single tensor or a list.
  // * "fun_in:0" gives the first element of a function input arg (a
  //   non-list input is considered a list of length 1 for these
  //   purposes).
  // * "node:out" where "node" is the name of a node in `node_def` and
  //   "out" is the name one of its op's output arguments (the name
  //   comes from the OpDef of the node's op). This represents that
  //   node's output, whether it is a single tensor or a list.
  //   Note: We enforce that an op's output arguments are never
  //   renamed in the backwards-compatibility test.
  // * "node:out:0" gives the first element of a node output arg (a
  //   non-list output is considered a list of length 1 for these
  //   purposes).
  //
  // NOT CURRENTLY SUPPORTED (but may be in the future):
  // * "node:out:-1" gives last element in a node output list
  // * "node:out:1:" gives a list with all but the first element in a
  //   node output list
  // * "node:out::-1" gives a list with all but the last element in a
  //   node output list

  // The body of the function.  Unlike the NodeDefs in a GraphDef, attrs
  // may have values of type `placeholder` and the `input` field uses
  // the "output" format above.
  repeated NodeDef node_def = 3;

  // A mapping from the output arg names from `signature` to the
  // outputs from `node_def` that should be returned by the function.
  map<string, string> ret = 4;
}

// GradientDef defines the gradient function of a function defined in
// a function library.
//
// A gradient function g (specified by gradient_func) for a function f
// (specified by function_name) must follow the following:
//
// The function 'f' must be a numerical function which takes N inputs
// and produces M outputs. Its gradient function 'g', which is a
// function taking N + M inputs and produces N outputs.
//
// I.e. if we have
//    (y1, y2, ..., y_M) = f(x1, x2, ..., x_N),
// then, g is
//    (dL/dx1, dL/dx2, ..., dL/dx_N) = g(x1, x2, ..., x_N,
//                                      dL/dy1, dL/dy2, ..., dL/dy_M),
// where L is a scalar-value function of (x1, x2, ..., xN) (e.g., the
// loss function). dL/dx_i is the partial derivative of L with respect
// to x_i.
message GradientDef {
  string function_name = 1;  // The function name.
  string gradient_func = 2;  // The gradient function's name.
}
