Fix perserving strides, inverse_indices in ManagedCollisionCollection #3597
+216
−0
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.
Summary:
The forward method of
ManagedCollisionCollectiongoes through each table, and maps the original indices of the input KJT into the new indices using a hash-function. This produces a Dict[str, JaggedTensor], which is then converted into a KeyedJaggedTensor.MCC should only change the values attribute of the KJT, while perserving all other attributes.
This conversion did not perserve key attributes of KJT such as
inverse_indices, andstride. that are essential to work with VBE.Reviewed By: kausv
Differential Revision: D84944895