Create a util function BuildGetTupleElementsForTupleResults
with proper processing on the shardings.
#64462
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.
Create a util function
BuildGetTupleElementsForTupleResults
with proper processing on the shardings.Issue 1.
The following code snippet is widely used in
mlir_hlo_to_hlo.cc
when exporting to a XlaOp if the op has multiple results, such as RngBitGeneratorOp and OptimizationBarrierOp. We create a util function to replace the repeating code snippets.Issue 2.
When we build
xla::GetTupleElement
in the code snippet, we inherit the tuple sharding (if any) and assign the tuple sharding to thexla::GetTupleElement
. This triggers a mismatching error on shape and sharding in XlaBuilder. This cl usesXlaScopedShardingAssignment
to process it appropriately.Enhance the test cases accordingly.