import tensorflow as tf

  def state_update_aggregation_modes(self):
    """See base class."""
    return {
        self.NORM_STATE_UPDATE_KEY: encoding_stage.StateAggregationMode.STACK
    }

  def initial_state(self):
    """See base class."""
    return {self.FACTOR_STATE_KEY: tf.constant(1.0)}

  # pylint: disable=g-doc-args,g-doc-return-or-yield
  def update_state(self, state, state_update_tensors):
    """Updates the state (see base class).

    This method illustrates how the implementation can handle state update based
    on a single encoding, or based on a multiple encodings collectively.
