schema.fbs

// Flatbuffers schema for serialized RTen models.
//
// See https://google.github.io/flatbuffers/flatbuffers_guide_writing_schema.html.
//
// Operator names and attributes align with the corresponding ONNX operators.
// See https://onnx.ai/onnx/operators/index.html.

file_identifier "RTEN";
file_extension "rten";

// Type for an operator node
enum OperatorType: ubyte {
  Add,
  ArgMin,
  ArgMax,
  AveragePool,
  BatchNormalization,
  Cast,
  Clip,
  Concat,
  ConstantOfShape,
  Conv,
  ConvTranspose,
  Cos,
  CumSum,
  Div,
  Equal,
  Erf,
  Expand,
  Flatten,
  Gather,
  Gemm,
  GlobalAveragePool,
  Greater,
  GRU,
  Identity,
  LeakyRelu,
  Less,
  LessOrEqual,
  Log,
  LogSoftmax,
  LSTM,
  MatMul,
  MaxPool,
  Mod,
  Mul,
  Pad,
  Pow,
  Range,
  ReduceMean,
  ReduceL2,
  Relu,
  Reshape,
  Resize,
  Shape,
  Sigmoid,
  Sin,
  Slice,
  Split,
  Sqrt,
  Squeeze,
  Softmax,
  Sub,
  Tanh,
  Transpose,
  Unsqueeze,
  Where,

  // New operators. These are appended here to preserve binary compatibility
  // of existing model files.
  ReduceProd,
  ReduceSum,
  ReduceMin,
  ReduceMax,
  NonZero,
  ScatterElements,
  Tile,
  Not,
  Abs,
  Max,
  Mean,
  Min,
  Sum,
  OneHot,
  Round,
  Floor,
  Ceil,
  Reciprocal,
  TopK,
  Neg,
  Exp,
  GreaterOrEqual,
  Size,
  Tan,
  Acos,
  Asin,
  Atan,
  InstanceNormalization,
  HardSigmoid,
  HardSwish,
  And,
  Or,
  Xor,
  Trilu,
  ScatterND,
  NonMaxSuppression,
  Sign,
  GatherElements,
  LayerNormalization,
  ReduceSumSquare,
  RandomUniform,
  Elu,
  RandomUniformLike,
  RandomNormal,
  RandomNormalLike,
  Softplus,
  GatherND,
}

enum RNNDirection: ubyte {
  Forward,
  Reverse,
  Bidirectional
}

enum AutoPad: ubyte {
  // nb. ONNX defines `SAME_UPPER` and `SAME_LOWER`. This corresponds to
  // `SAME_UPPER` and TensorFlow / Keras's "same".
  Same,
  NotSet
}

enum DataType: ubyte {
  Int32,
  Float
}

// Coordinate transform modes for Resize operator.
enum CoordTransformMode: ubyte {
  HalfPixel,
  Asymmetric,
  AlignCorners
}

// Rounding modes supported by Resize operator when `ResizeMode` is `Nearest`.
enum NearestMode: ubyte {
  Floor,
  Ceil,
  RoundPreferFloor,
  RoundPreferCeil,
}

enum ResizeMode: ubyte {
  Nearest,
  Linear
}

// Operator-specific configuration
union OperatorAttrs {
  ArgMaxAttrs, // Also used for ArgMin
  AveragePoolAttrs,
  BatchNormalizationAttrs, // Also used for InstanceNormalization
  CastAttrs,
  ConcatAttrs,
  ConstantOfShapeAttrs,
  ConvAttrs,
  ConvTransposeAttrs,
  FlattenAttrs,
  GatherAttrs, // Also used for GatherElements
  GemmAttrs,
  GRUAttrs,
  LeakyReluAttrs,
  LSTMAttrs,
  MaxPoolAttrs,
  ReduceMeanAttrs, // Also used for Reduce* ops
  ReshapeAttrs,
  ResizeAttrs,
  SplitAttrs,
  SoftmaxAttrs, // Also used for LogSoftmax
  TransposeAttrs,

  // New operator attrs. These are appended here to preserve binary compatibility
  // of existing model files.
  ModAttrs,
  ScatterElementsAttrs,
  OneHotAttrs,
  TopKAttrs,
  HardSigmoidAttrs,
  TriluAttrs,
  ScatterNDAttrs,
  NonMaxSuppressionAttrs,
  LayerNormalizationAttrs,
  RandomUniformAttrs,
  EluAttrs,
  RandomUniformLikeAttrs,
  RandomNormalAttrs,
  RandomNormalLikeAttrs,
  GatherNDAttrs,
}

table ArgMaxAttrs {
  axis:int;
  keep_dims:bool;
}

table AveragePoolAttrs {
  kernel_size:[uint] (required);
  auto_pad:AutoPad;

  // Padding for spatial axes as [top, left, bottom, right]
  pads:[uint];

  strides:[uint];

  count_include_pad:bool;
}

table BatchNormalizationAttrs {
  epsilon:float;
}

table CastAttrs {
  to:DataType;
}

table ConcatAttrs {
  axis:int;
}

union Scalar {
  IntScalar,
  FloatScalar
}

table IntScalar {
  value:int32;
}

table FloatScalar {
  value:float32;
}

table ConstantOfShapeAttrs {
  value:Scalar (required);
}

table ConvAttrs {
  auto_pad:AutoPad;

  // Padding for spatial axes as [left, right] or [top, left, bottom, right]
  pads:[uint];

  groups:uint;
  strides:[uint];
  dilations:[uint];
}

table ConvTransposeAttrs {
  strides:[uint];

  // Defaults to `NotSet` for backwards compatibility.
  auto_pad:AutoPad = NotSet;

  // Padding for spatial axes as [left, right] or [top, left, bottom, right]
  pads:[uint];
}

table EluAttrs {
  alpha:float;
}

table FlattenAttrs {
  axis:int;
}

table LayerNormalizationAttrs {
  axis:int;
  epsilon:float;
}

table GatherAttrs {
  axis:int;
}

table GatherNDAttrs {
  batch_dims:int;
}

table GemmAttrs {
  alpha:float;
  beta:float;
  transpose_a:bool;
  transpose_b:bool;
}

table GRUAttrs {
  direction:RNNDirection;
  hidden_size:uint;
  linear_before_reset:bool;
}

table HardSigmoidAttrs {
  alpha:float;
  beta:float;
}

table LeakyReluAttrs {
  alpha:float;
}

table LSTMAttrs {
  direction:RNNDirection;
  hidden_size:uint;
}

table MaxPoolAttrs {
  kernel_size:[uint] (required);
  auto_pad:AutoPad;

  // Padding for spatial axes as [top, left, bottom, right]
  pads:[uint];

  strides:[uint];
}

table ModAttrs {
  fmod:bool;
}

enum NMSBoxOrder: ubyte {
  TopLeftBottomRight,
  CenterWidthHeight,
}

table NonMaxSuppressionAttrs {
  box_order:NMSBoxOrder;
}

table OneHotAttrs {
  axis:int;
}

table RandomNormalAttrs {
  mean:float;
  scale:float;
  seed:float = null;
  shape:[uint];
}

table RandomNormalLikeAttrs {
  mean:float;
  scale:float;
  seed:float = null;
}

table RandomUniformAttrs {
  shape:[uint];
  high:float;
  low:float;
  seed:float = null;
}

table RandomUniformLikeAttrs {
  high:float;
  low:float;
  seed:float = null;
}

table ReduceMeanAttrs {
  axes:[int];
  keep_dims:bool;
}

table ReshapeAttrs {
  allow_zero:bool;
}

table ResizeAttrs {
  mode:ResizeMode;
  coord_mode:CoordTransformMode;
  nearest_mode:NearestMode;
}

enum ScatterReduction: ubyte {
  None,
  Add,
  Mul,
  Min,
  Max,
}

table ScatterElementsAttrs {
  axis:int;
  reduction:ScatterReduction;
}

table ScatterNDAttrs {
  reduction:ScatterReduction;
}

table SoftmaxAttrs {
  axis:int;
}

table SplitAttrs {
  axis:int;
}

table TopKAttrs {
  axis:int;
  largest:bool;
  sorted:bool;
}

table TransposeAttrs {
  perm:[uint];
}

table TriluAttrs {
  upper:bool;
}

// Node in the dataflow graph
union NodeKind {
  OperatorNode,
  ConstantNode,
  ValueNode
}

// Graph node that computes an output tensor given one or more inputs and
// operator configuration.
table OperatorNode {
  type:OperatorType;
  attrs:OperatorAttrs;

  // Indexes of input nodes. Negative values indicate missing optional inputs.
  inputs:[int];

  // Indexes of output nodes. Negative values indicate unused outputs.
  outputs:[int];
}

union ConstantData {
  FloatData,
  IntData,
}

table FloatData {
  data: [float32] (required);
}

table IntData {
  data: [int32] (required);
}

// Graph node for a constant tensor value, whose data is part of the model.
table ConstantNode {
  shape:[uint] (required);
  data:ConstantData (required);
}

// Dimension of a ValueNode's shape. This can be either a fixed value or a
// symbolic name.
//
// This is logically a union, but uses a table due to
// https://github.com/google/flatbuffers/issues/5024.
table Dim {
  value:uint;
  name:string;
}

// Graph node for a dynamic tensor value, such as a model input or operator
// output.
table ValueNode {
  // Expected shape of the tensor at runtime.
  shape:[Dim];
}

table Node {
  // Identifier for external referencing and debugging
  name:string;
  data:NodeKind;
}

// Dataflow graph for the model
table Graph {
  // Nodes are sorted in topological order. This means that if a node references
  // other nodes (eg. an operator node referencing inputs), the referents must
  // come earlier in the list.
  nodes:[Node];

  // IDs of input nodes
  inputs:[uint];

  // IDs of output nodes
  outputs:[uint];
}

table Metadata {
  // SHA-256 hash of the ONNX model that was used as the source for this RTen
  // model.
  onnx_hash:string;

  // A short description of what this model does.
  description:string;

  // Identifier for the license used in this model.
  //
  // This should be an SPDX (https://spdx.org/licenses/) identifier for openly
  // licensed models.
  license:string;

  // Commit ID for the code that produced this model.
  commit:string;

  // URL of repository where the model's code is hosted (eg. GitHub).
  code_repository:string;

  // URL of repository where the model is hosted (eg. Hugging Face).
  model_repository:string;

  // Identifier for the training run that produced this model.
  run_id:string;

  // URL of logs etc. for the training run that produced this model.
  run_url:string;
}

table Model {
  schema_version:int;
  graph:Graph (required);
  metadata:Metadata;
}

root_type Model;