docs: def initial_state(self, batch_size, trainable=False):
      """Creates the initial memory.
  
      We should ensure each row of the memory is initialized to be unique,
      so initialize the matrix to be the identity. We then pad or truncate
      as necessary so that init_state is of size
      (batch_size, self._mem_slots, self._mem_size).
  
      Args:
        batch_size: The size of the batch.
        trainable: Whether the initial state is trainable. This is always True.
  
      Returns:
        init_state: A truncated or padded matrix of size
          (batch_size, self._mem_slots, self._mem_size).
      """
      init_state = tf.eye(self._mem_slots, batch_shape=[batch_size])
  
      # Pad the matrix with zeros.
      if self._mem_size > self._mem_slots:
        difference = self._mem_size - self._mem_slots
        pad = tf.zeros((batch_size, self._mem_slots, difference))
        init_state = tf.concat([init_state, pad], -1)
      # Truncation. Take the first `self._mem_size` components.
      elif self._mem_size < self._mem_slots:
        init_state = init_state[:, :, :self._mem_size]
      return init_state