docs: def _build(self):
      """Connects the module to the graph.
  
      Returns:
        The learnable state, which has the same type, structure and shape as
          the `initial_state` passed to the constructor.
      """
      flat_initial_state = nest.flatten(self._initial_state)
      if self._mask is not None:
        flat_mask = nest.flatten(self._mask)
        flat_learnable_state = [
            _single_learnable_state(state, state_id=i, learnable=mask)
            for i, (state, mask) in enumerate(zip(flat_initial_state, flat_mask))]
      else:
        flat_learnable_state = [_single_learnable_state(state, state_id=i)
                                for i, state in enumerate(flat_initial_state)]
  
      return nest.pack_sequence_as(structure=self._initial_state,
                                   flat_sequence=flat_learnable_state)