AdaptiveEncodingStageInterface decode should be able to receive and update a state

System information

• TensorFlow version (you are using): 2.2.1
• Are you willing to contribute it (Yes/No): Maybe

Motivation

What is the use case and how does it broadly benefits users? Prioritization and whether a feature is added is based on how it
helps the community and the feature's maintenance costs.

The benefit is more communication efficient ML and more consolidated research in that field.

As examples:

Some encoding processes maintain a server state, such as server-side error-feedback and pre-decode momentum

For example, see https://arxiv.org/abs/1802.04434 (signum) https://arxiv.org/abs/2007.07682 (fetchsgd)

Describe the feature
currently only encode is able to update the encoding-stage state.

it would be beneficial if decode could also return updated-state-tensors (which would do not need to be aggregated)

for example, pre-decode momentum tensor would be updated in this way (momentum on the possibly aggregated encoded values)

Start by locating AdaptiveEncodingStageInterface and comparing how encode and decode handle state. Define the requested decode behavior so updated-state tensors are returned without aggregation, then verify it against the server-side state use cases described in the issue.