syntax = "proto3";

package tensorflow.ifrt_serving;

import "xla/xla_data.proto";
import "tensorflow/core/framework/tensor.proto";

enum IfrtPjRtServingPlatformType {
  IFRT_PJRT_SERVING_PLATFORM_TYPE_UNSPECIFIED = 0;
  IFRT_PJRT_SERVING_PLATFORM_TYPE_TPU = 1;
  IFRT_PJRT_SERVING_PLATFORM_TYPE_GPU = 2;
}

// `tf.IfrtLoadVariableOp`'s configuration message for device and spmd
// configuration.
message VariableDeviceShardingConfigProto {
  xla.OpSharding sharding = 1;
  repeated int32 device_ids = 2;
}

enum IfrtServingCoreSelectionPolicy {
  // Default policy to select the soonest-to-finish core defined at
  // http://shortn/_XQhI1ASAu0.
  IFRT_SERVING_CORE_SELECTION_POLICY_DEFAULT = 0;
  // Policy that round robin with local ordinal http://shortn/_7BtVe4dkp5.
  IFRT_SERVING_CORE_SELECTION_POLICY_LOCAL_ROUND_ROBIN = 1;
}

message DefaultSignatureInputConfig {
  message Signature {
    map<string, tensorflow.TensorProto> default_inputs = 1;
  }

  map<string, Signature> signatures = 1;
}
