aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
diff options
context:
space:
mode:
authorGravatar Mark Heffernan <meheff@google.com>2018-07-03 10:34:20 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-07-03 10:37:09 -0700
commit31a5fa1ee88f8f3bb1a46f3734136b6d85e8642f (patch)
tree8ada804c9a30d162778decc9fcef009acfa39d49 /tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
parentbce9c2ef1ad72a5d962faac0d114932af6a69bf9 (diff)
Change Send, SendDone, Recv and RecvDone to produce tokens.
This is a follow up to cl/202069017 which added tokens as operands to Send and Recv. PiperOrigin-RevId: 203145403
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_module_group_metadata.cc')
-rw-r--r--tensorflow/compiler/xla/service/hlo_module_group_metadata.cc3
1 files changed, 2 insertions, 1 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
index bf33640db1..6bcd7b042d 100644
--- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
+++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
@@ -382,7 +382,8 @@ Status HloModuleGroupMetadata::VerifyChannelInstructions() {
// Check if the shapes match for each channel.
for (const Channel& channel : channels_) {
const Shape& send_shape = channel.send->operand(0)->shape();
- const Shape& recv_shape = channel.recv_done->shape();
+ const Shape& recv_shape =
+ ShapeUtil::GetTupleElementShape(channel.recv_done->shape(), 0);
if (!ShapeUtil::Compatible(send_shape, recv_shape)) {
return FailedPrecondition("send/recv shapes do not match");
}