Const fold scatter_along_axis and scaled_dot_product_attention correctly - #2802
Open
LeSingh1 wants to merge 1 commit into
Open
Const fold scatter_along_axis and scaled_dot_product_attention correctly#2802LeSingh1 wants to merge 1 commit into
LeSingh1 wants to merge 1 commit into
Conversation
Both ops had a value_inference that ignored an input which changes the result, so a program whose operands happen to be const folds to the wrong value and the op disappears from the model. scatter_along_axis: value_inference always called np.put_along_axis, i.e. it always computed mode="update". Every other mode was folded to the overwrite result, including the default mode="add". Compute the mode's reduction with the matching numpy ufunc applied through ufunc.at so that repeated indices accumulate, which is what the runtime does. scaled_dot_product_attention: value_inference only applied attn_mask when the mask had a value, and otherwise silently produced unmasked attention. This op has no @precondition, so with const query/key/value and a mask computed at runtime the whole attention collapsed to a const of the unmasked result. Return None instead, so the op stays in the graph.
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
1 participant
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.
Two
value_inferenceimplementations ignore an input that changes the result, so const folding produces a different answer than execution.scatter_along_axisalways computesmode="update", but the op's default isadd.scaled_dot_product_attentionignoresattn_mask, so a const-folded call silently produces unmasked attention.Both now honour the input, or decline to fold when they can't.
Tested with real predictions on this machine: each new test fails on
mainand passes here.