diff options
Diffstat (limited to 'tensorflow/core/kernels/roll_op.cc')
| -rw-r--r-- | tensorflow/core/kernels/roll_op.cc | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/tensorflow/core/kernels/roll_op.cc b/tensorflow/core/kernels/roll_op.cc index 722116f86f..efa30438d9 100644 --- a/tensorflow/core/kernels/roll_op.cc +++ b/tensorflow/core/kernels/roll_op.cc @@ -19,6 +19,7 @@ limitations under the License. #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/register_types_traits.h" #include "tensorflow/core/framework/shape_inference.h" +#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/types.h" #include "tensorflow/core/util/work_sharder.h" @@ -258,7 +259,7 @@ class RollOp : public OpKernel { if (axis < 0) { axis += num_dims; } - OP_REQUIRES(context, 0 <= axis && axis < num_dims, + OP_REQUIRES(context, FastBoundsCheck(axis, num_dims), errors::InvalidArgument("axis ", axis, " is out of range")); const int ds = std::max<int>(static_cast<int>(input.dim_size(axis)), 1); const int sum = shift_mod_sum[axis] + static_cast<int>(shift_flat(i)); |