import tensorflow as tf

    return {self.ENCODED_VALUES_KEY: x + addend}

  def decode(self,
             encoded_tensors,
             decode_params,
             num_summands=None,
             shape=None):
    """See base class."""
    del num_summands  # Unused.
    del shape  # Unused.
    x = encoded_tensors[self.ENCODED_VALUES_KEY]
    addend = dummy_rng_source(decode_params[self.SEED_PARAM_KEY],
                              x.shape.num_elements())
    addend = tf.reshape(addend, x.shape)
    return x - addend


@encoding_stage.tf_style_adaptive_encoding_stage
class PlusOneOverNEncodingStage(encoding_stage.AdaptiveEncodingStageInterface):
  """[Example] adaptive encoding stage, adding 1/N in N-th iteration.

  This is an example implementation of an `AdaptiveEncodingStageInterface` that
  modifies state, which controls the creation of params. This is also a simple
  example of how an `EncodingStageInterface` can be wrapped as an
  `AdaptiveEncodingStageInterface`, without modifying the wrapped encode and
  decode methods.
