[xla] Change Shape::IsArray and Shape::IsArrayOrBuffer. #97389
Merged
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.
[xla] Change Shape::IsArray and Shape::IsArrayOrBuffer.
Currently, Shape::IsArray returns true when the shape is an array primitive and Shape::IsArrayOrBuffer returns true when the shape is an array primitive or a buffer primitive. The majority use of Shape::IsArray is to exclude TOKEN/TUPLE so that we can further find out the storage element_type for the Shape. For this reason, In order to support the new Buffer type, we would need to change most of uses of Shape::IsArray to Shape::IsArrayOrBuffer, leaving only a few uses of Shape::IsArray within shape_util and shape_inference.
To avoid such a massive change, we decide to make Shape::IsArray to return true when the shape is either an array or a buffer, which is what the current Shape::IsArrayOrBuffer does. We also move the current Shape::IsArray to a new routine Shape::IsArrayExcludingBuffer, which should have very limited uses, such as in shape_inference and shape_util.