import tensorflow as tf

    new_factor = (
        weight * state[self.FACTOR_STATE_KEY] + (1 - weight) / norm_mean)
    return {self.FACTOR_STATE_KEY: new_factor}

  def get_params(self, state):
    """See base class."""
    params = {self.FACTOR_PARAM_KEY: state[self.FACTOR_STATE_KEY]}
    return params, params

  def encode(self, x, encode_params):
    """See base class."""
    return (self._stage.encode(x, encode_params), {
        self.NORM_STATE_UPDATE_KEY: tf.norm(x)
    })

  def decode(self,
             encoded_tensors,
             decode_params,
             num_summands=None,
             shape=None):
    """See base class."""
    return self._stage.decode(encoded_tensors, decode_params, num_summands,
                              shape)

