From 1dbc6712045108d0d50f6a3b7d5a749322b6843a Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Tue, 17 Apr 2018 01:14:03 +0000 Subject: [PATCH] Check in shape function that axis and shift are same size Signed-off-by: Yong Tang --- tensorflow/core/ops/manip_ops.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/ops/manip_ops.cc b/tensorflow/core/ops/manip_ops.cc index 8461b1d..c90b2b2 100644 --- a/tensorflow/core/ops/manip_ops.cc +++ b/tensorflow/core/ops/manip_ops.cc @@ -36,7 +36,8 @@ REGISTER_OP("Roll") TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 1, &unused)); // The `axis` must be scalar or 1-D. TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(2), 1, &unused)); - + // Validate 'shift' is the same shape as axis'. + TF_RETURN_IF_ERROR(c->Merge(c->input(1), c->input(2), &unused) ); return shape_inference::UnchangedShape(c); }); -- 2.7.4