Repository navigation
check element dtype against parent in batch_util ValidateInput - #130332
Open
madib06ops wants to merge 1 commit into
Open
madib06ops wants to merge 1 commit into
madib06ops wants to merge 1 commit into
Conversation
CopyElementToSlice, CopySliceToElement and MaybeMoveSliceToElement only compared element counts before indexing the parent with base<T>() of one side's dtype, so an element whose dtype differs from the parent was copied with the wrong element size, past the end of the parent buffer. Reject mismatched dtypes the same way CopyContiguousSlices and MaybeMoveContiguousSlices already do.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a data type validation check in ValidateInput within tensorflow/core/util/batch_util.cc to ensure that the element and parent tensors have matching data types, returning an absl::InternalError if they do not. It also adds corresponding unit tests in tensorflow/core/framework/batch_util_test.cc to verify this behavior for CopyElementToSlice, CopySliceToElement, and MaybeMoveSliceToElement. There are no review comments to evaluate, and the implementation is correct and well-tested, so I have no feedback to provide.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
batch_util::ValidateInput only compares element counts, so CopyElementToSlice, CopySliceToElement and MaybeMoveSliceToElement copy a slice using one tensor's dtype while indexing the other's buffer; when a tf.data element's dtype differs from the batch tensor (for example a crafted SnapshotChunkDataset file whose records are parsed with Tensor::FromProto and never checked against output_types), Batch writes num_values*sizeof(T) bytes at the wrong stride past the end of the batch buffer, and on TF 2.20 an int8 record followed by an int64 record of 65536 elements batched by 2 corrupts the heap and segfaults. This rejects mismatched dtypes in ValidateInput, matching the check CopyContiguousSlices and MaybeMoveContiguousSlices already do, and adds DifferentDtypes tests to batch_util_test.cc that fail before the change and pass after.