syntax = "proto3";

package syft.lib.torch;

import "proto/lib/torch/device.proto";

message TensorData {
  repeated int64 shape = 1;
  string dtype = 2;

  bool is_quantized = 3;
  float scale = 4;
  int32 zero_point = 5;

  // Field numbers starting at 16 take two bytes to encode,
  // so starting the tensor data at 16 leaves room for more
  // commonly occurring fields to have one byte field numbers

  // Only the contents field corresponding to the value of the
  // dtype field should be set, so most of the contents fields
  // will not be present in any given TensorData

  // Integer types
  repeated uint32 contents_uint8 = 16;
  repeated int32 contents_int8 = 17;
  repeated int32 contents_int16 = 18;
  repeated int32 contents_int32 = 19;
  repeated int64 contents_int64 = 20;

  // Floating point types
  repeated float contents_float16 = 21;
  repeated float contents_float32 = 22;
  repeated double contents_float64 = 23;

  // Boolean types
  repeated bool contents_bool = 24;

  // Quantized types
  repeated sint32 contents_qint8 = 25;  // 8-bit signed integer
  repeated uint32 contents_quint8 = 26; // 8-bit unsigned integer
  repeated sint32 contents_qint32 = 27; // 32-bit signed integer

  // Brain floating point types
  repeated float contents_bfloat16 = 28; // float32 truncated to first 16 bits

  // Complex types (not yet implemented in PyTorch)
  // repeated float contents_complex32 = 29;
  // repeated float contents_complex64 = 30;
  // repeated double contents_complex128 = 31;
}

message TensorProto {
  TensorData tensor = 1;
  bool requires_grad = 2;
  TensorData grad = 3;
  syft.lib.torch.Device device = 4;
}
