Function IsMultiDevice

tensorflow/core/common_runtime/eager/execute.cc:321–345

}

bool IsMultiDevice(const FunctionDef* fdef) {
  if (fdef == nullptr) {
    // Primitive op.
    return false;
  }

  // Run all functions as multi-device.
  return true;

  // We can eliminate some overhead by running simple functions using regular
  // CallOp kernel. However, it is tricky to figure out which functions should
  // be run using CallOp. Also, currently CallOp runs neither optimization
  // passes (needed for TPU/XLA) nor grappler.
  // Here are some cases where a function should be run in multi-device mode:
  //  - Function takes at least two resources on different devices.
  //  - Function takes a resource on deviceA and a body op explicitly placed
  //  on deviceB.
  //  - Function has a colocation constraint.
  //  - Function has an explicit device annotation (which might not be using
  //    full canonical device name) different from op_device. Note that false
  //    positives are ok.
  //  - Function has a node or a (node) attribute that can potentially make
  //    the function multi-device after a rewrite pass (e.g. various XLA/TPU
  //    special nodes and attributes)
}

Status GetDeviceForInput(const EagerContext* ctx, TensorHandle* tensor_handle,
                         Device** result) {