message: Avoid redundant function input copies for int32 on TPUs.
diff: @@ -302,13 +302,20 @@ Status GetDeviceForInput(const EagerContext& ctx, TensorHandle* tensor_handle,
  TF_RETURN_IF_ERROR(
  ctx.FindDeviceFromName(device_name.c_str(), &input_device));
  *result = input_device;
  - } else if (MTypeFromDType(tensor_handle->dtype) == HOST_MEMORY) {
  - *result = cpu_device;
  } else {
  Device* device = absl::get<Device*>(tensor_handle->device());
  + const bool is_tpu = device != nullptr && device->device_type() == "TPU";
  + // int32 return values can be placed on TPUs.
  + const bool use_host_memory =
  + is_tpu ? MTypeFromDTypeIntsOnDevice(tensor_handle->dtype)
  + : MTypeFromDType(tensor_handle->dtype);
  + if (use_host_memory) {
  + *result = cpu_device;
  + } else {
  device_name = device != nullptr ? device->name() : cpu_device->name();
  *result = (device == nullptr ? cpu_device : device);
  }
  + }
  return Status::OK();
  }
  