Add inputs_k
and inputs_v
args to attention layer
#3379
Closed
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.
Currently,
MultiHeadDotProductAttention
layer's call method signature isMultiHeadDotProductAttention.__call__(inputs_q, inputs_kv, mask=None, deterministic=None)
. As discussed in #1737, there are some cases where passing in separate values for the key and values is desired, which isn't possible with the current API. This PR adds two more arguments,inputs_k
andinputs_v
to the call method signature and sets the method signature to the following:MultiHeadDotProductAttention.__call__(inputs_q, inputs_k=None, inputs_v=None, *, inputs_kv=None, mask=None, deterministic=None)
. Note that theinputs_kv
,mask
anddeterministic
args are now keyword arguments.inputs_k
andinputs_v
areNone
, then they will both copy the value ofinputs_q
(i.e. self attention)inputs_v
isNone
, it will copy the value ofinputs_k
(same behavior as the previous API, i.e.module.apply(inputs_q=query, inputs_k=key_value, ...)
is equivalent tomodule.apply(inputs_q=query, inputs_kv=key_value, ...)
)inputs_kv
is not None, bothinputs_k
andinputs_v
will copy the value ofinputs_kv
Users can still use
inputs_kv
but aDeprecationWarning
will be raised andinputs_kv
will be removed in the future.Since self attention can be done using this new API, the
SelfAttention
layer will also raise aDeprecationWarning
and will be removed in the future.Check out #3389 to see examples of how to port your code over to the new API.