1 files changed, 15 insertions, 1 deletions
diff --git a/tensorflow/compiler/xla/client/client.cc b/tensorflow/compiler/xla/client/client.cc
index 3a157c69cd..8e54311bad 100644
--- a/tensorflow/compiler/xla/client/client.cc
+++ b/tensorflow/compiler/xla/client/client.cc
@@ -409,8 +409,10 @@ StatusOr<string> Client::ExecutionStatsAsString(
   return string("[Execution Statistics] not available.");
 }
 
-StatusOr<ChannelHandle> Client::CreateChannelHandle() {
+StatusOr<ChannelHandle> Client::CreateChannelHandleByType(
+    ChannelHandle::ChannelType type) {
   CreateChannelHandleRequest request;
+  request.set_channel_type(type);
   CreateChannelHandleResponse response;
 
   VLOG(1) << "making create channel handle request";
@@ -424,4 +426,16 @@ StatusOr<ChannelHandle> Client::CreateChannelHandle() {
   return response.channel();
 }
 
+StatusOr<ChannelHandle> Client::CreateChannelHandle() {
+  return CreateChannelHandleByType(ChannelHandle::DEVICE_TO_DEVICE);
+}
+
+StatusOr<ChannelHandle> Client::CreateHostToDeviceChannelHandle() {
+  return CreateChannelHandleByType(ChannelHandle::HOST_TO_DEVICE);
+}
+
+StatusOr<ChannelHandle> Client::CreateDeviceToHostChannelHandle() {
+  return CreateChannelHandleByType(ChannelHandle::DEVICE_TO_HOST);
+}
+
 }  // namespace xla