From 729192823935156ae29d7f0d5f64c0bcd6034c7a Mon Sep 17 00:00:00 2001 From: Jacques Pienaar Date: Fri, 20 Apr 2018 10:32:24 -0700 Subject: [PATCH] Adding Shape inference functions to outfeed enqueue ops. PiperOrigin-RevId: 193688099 --- tensorflow/contrib/tpu/ops/outfeed_ops.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorflow/contrib/tpu/ops/outfeed_ops.cc b/tensorflow/contrib/tpu/ops/outfeed_ops.cc index 5900c61..b05c76c 100644 --- a/tensorflow/contrib/tpu/ops/outfeed_ops.cc +++ b/tensorflow/contrib/tpu/ops/outfeed_ops.cc @@ -26,6 +26,7 @@ REGISTER_OP("OutfeedEnqueue") .Input("input: dtype") .Attr("dtype: type") .SetIsStateful() + .SetShapeFn(shape_inference::NoOutputs) .Doc(R"doc( An op which emits a single Tensor value from an XLA computation. @@ -36,6 +37,7 @@ REGISTER_OP("OutfeedEnqueueTuple") .Input("inputs: dtypes") .Attr("dtypes: list(type)") .SetIsStateful() + .SetShapeFn(shape_inference::NoOutputs) .Doc(R"doc( An op which emits multiple Tensor values from an XLA computation. -- 2.7.4