syntax = "proto2";

import "dragnn/protos/spec.proto";

package syntaxnet.dragnn.runtime;

// Specification of a subgraph of TF nodes that make up a network cell.
//
// Roughly speaking, a "cell" consists of the "pure math" parts of a DRAGNN
// component, and is intended to be exported to a NN compiler.  The set of
// operations that make up a cell may change over time, but currently the
// boundaries of a cell are:
//
// Inputs:
//   * Fixed feature IDs.
//   * Linked feature embeddings, before pass_through_embedding_matrix().
//   * Recurrent context tensors.
//
// Outputs:
//   * Network unit layers.
message CellSubgraphSpec {
  // An input to the subgraph.
  message Input {
    // Possible types of input.
    enum Type {
      TYPE_UNKNOWN = 0;

      // An input derived from a fixed or linked feature.
      TYPE_FEATURE = 1;

      // An input that refers to an output of the previous iteration of the
      // transition loop.  The input must have the same name as the output to
      // which it refers.  On the first iteration, its value is zero.
      //
      // This is used by, e.g., LSTMNetwork, which reads its cell state from the
      // context_tensor_arrays instead of from a linked feature.
      TYPE_RECURRENT = 2;
    }

    // Logical name of the input (e.g., "lstm_c", "linked_feature_0").  Must be
    // unique among the inputs of the cell.
    optional string name = 1;

    // Tensor containing the input (e.g., "annotation/rnn/split:1").  Must be
    // unique among the inputs of the cell.
    optional string tensor = 2;

    // Type of input.
    optional Type type = 3 [default = TYPE_UNKNOWN];
  }

  // An output of the subgraph.
  message Output {
    // Logical name of the output (e.g., "lstm_c", "layer_0").  Must be unique
    // among the outputs of the cell.
    optional string name = 1;

    // Tensor containing the output (e.g., "annotation/rnn/split:1").  Need not
    // be unique; duplicate outputs for the same tensor are treated as aliases.
    optional string tensor = 2;
  }

  // Inputs of the subgraph.
  repeated Input input = 1;

  // Outputs of the subgraph.
  repeated Output output = 2;
}

// Additional information to compile a component.
//
// NEXT ID: 3
message CompilationSpec {
  extend ComponentSpec {
    optional CompilationSpec component_spec_extension = 174770970;
  }

  // A unique name of the entire DRAGNN model where this component is used.
  optional string model_name = 1;

  // The subgraph specification for this component.
  optional CellSubgraphSpec cell_subgraph_spec = 2;
}