Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix heap OOB in tf.raw_ops.RaggedGather
PiperOrigin-RevId: 388355464
Change-Id: If14d96231d1cd7aad7c4d1c22c1bab1576b75717
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Aug 3, 2021
1 parent 4979e3b commit a2b743f
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion tensorflow/core/kernels/ragged_gather_op.cc
Expand Up @@ -58,15 +58,21 @@ class RaggedGatherOpBase : public OpKernel {

void Compute(OpKernelContext* context) override {
// Get the input Tensors.

OpInputList params_nested_splits_in;
OP_REQUIRES_OK(context, context->input_list("params_nested_splits",
&params_nested_splits_in));
OP_REQUIRES(
context, params_nested_splits_in.size() > 0,
errors::InvalidArgument("params_nested_splits must be non empty"));

const Tensor& params_dense_values_in =
context->input(params_nested_splits_in.size());
const Tensor& indices_in =
context->input(params_nested_splits_in.size() + 1);

DCHECK_GT(params_nested_splits_in.size(), 0); // Enforced by REGISTER_OP.
OP_REQUIRES(context, params_nested_splits_in[0].dims() > 0,
errors::InvalidArgument("Split tensors must not be scalars"));
SPLITS_TYPE num_params = params_nested_splits_in[0].dim_size(0) - 1;
OP_REQUIRES_OK(context, ValidateIndices(indices_in, num_params));

Expand Down

0 comments on commit a2b743f

Please sign in to comment.