Skip to content
Permalink
Browse files Browse the repository at this point in the history
Allow 0 for number of segments in unsorted_segment_join_op.cc
Related to the fix for #55305

PiperOrigin-RevId: 443157549
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Apr 20, 2022
1 parent 3f30e49 commit 20cb187
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions tensorflow/core/kernels/unsorted_segment_join_op.cc
Expand Up @@ -94,8 +94,10 @@ class UnsortedSegmentJoinOp : public OpKernel {
errors::InvalidArgument("Number of segments cannot be empty."));
auto num_segments = num_segments_tensor.scalar<NUM_SEGMENTS_TYPE>()();

OP_REQUIRES(context, num_segments > 0,
errors::InvalidArgument("Number of segments must be positive"));
OP_REQUIRES(
context, num_segments >= 0,
errors::InvalidArgument(
"Number of segments must be non-negative but got ", num_segments));
OP_REQUIRES(context, segment_dims != 0,
errors::InvalidArgument("Segment_id cannot have rank 0"));

Expand Down

0 comments on commit 20cb187

Please sign in to comment.