docs: def state_size(self):
      """Tuple of `tf.TensorShape`s indicating the size of state tensors."""
      if self._max_unique_stats == 1:
        return (tf.TensorShape([self._hidden_size]),
                tf.TensorShape([self._hidden_size]))
      else:
        return (tf.TensorShape([self._hidden_size]),
                tf.TensorShape([self._hidden_size]),
                tf.TensorShape(1))