message: Allow the XLA device assignment to be passed to the XRT compilation machinery.
diff: @@ -6,6 +6,24 @@ import "tensorflow/compiler/tf2xla/host_compute_metadata.proto";
  import "tensorflow/compiler/xla/xla_data.proto";
  import "tensorflow/compiler/xla/service/hlo.proto";
  +message DeviceAssignment {
  + message ComputationDevice {
  + message DeviceMeshCoordinates {
  + // The mesh coordinates for the device. Usually (X, Y, Core), in the order
  + // in which they are returned in the TopologyProto.
  + // X = value(0)
  + // Y = value(1)
  + // Core = value(2)
  + repeated int32 value = 1;
  + }
  + // As many replicas as there are in the replicated computation.
  + repeated DeviceMeshCoordinates replica_devices = 1;
  + }
  + // As many ComputationDevice as many there are computations (number
  + // of cores per replica).
  + repeated ComputationDevice computation_devices = 1;
  +}
  +
  // Options for an XLA compilation.
  message XLAComputationConfig {
  // The number of replicas the computation will be run on. If this is
  @@ -23,6 +41,11 @@ message XLAComputationConfig {
  // computation. per_core_args_and_result_shapes is optional for a
  // single-core computation.
  repeated xla.ProgramShape per_core_program_shape = 5;
  + // Describes how replicated computation instances should be assigned to
  + // devices. There are num_cores_per_replica computations, and each one will be
  + // sent and executed to the set of replica device numbers described in the
  + // DeviceAssignment proto.
  + DeviceAssignment device_assignment = 6;
  }
  // Options and XLA computation for a compilation.
  