Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix heap-buffer-overflow issue with tf.raw_ops.SparseReshape.
PiperOrigin-RevId: 371218558
Change-Id: I6a6dc5bf15b50a1d05bdd95e9ba347cb39f40f45
  • Loading branch information
Amit Patankar authored and tensorflower-gardener committed Apr 29, 2021
1 parent 8d78df9 commit 1d04d7d
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions tensorflow/core/kernels/sparse_reshape_op.cc
Expand Up @@ -26,6 +26,7 @@ limitations under the License.
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/reshape_util.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/core/platform/errors.h"

namespace tensorflow {

Expand All @@ -38,6 +39,17 @@ class SparseReshapeOp : public OpKernel {
explicit SparseReshapeOp(OpKernelConstruction* context) : OpKernel(context) {}

void Compute(OpKernelContext* context) override {
const Tensor& input_indices_in = context->input(0);
const Tensor& input_shape_in = context->input(1);

OP_REQUIRES(context, TensorShapeUtils::IsMatrix(input_indices_in.shape()),
errors::InvalidArgument("Input must be a matrix."));
OP_REQUIRES(context, TensorShapeUtils::IsVector(input_shape_in.shape()),
errors::InvalidArgument("Input shape must be a vector."));
OP_REQUIRES(context,
input_indices_in.dim_size(1) == input_shape_in.dim_size(0),
errors::InvalidArgument(
"Input tensor rank must match input shape length."));
ReshapeSparseTensor<Device>(context, context->input(0), context->input(1),
context->input(2), 0 /* output indices index */,
1 /* output shape index */);
Expand Down

0 comments on commit 1d04d7d

Please sign in to comment.