def state(self, state):
    """Sets the state of this layer and its sublayers.

    Args:
      state: the state to set; if layer has sublayers, state should be
        either a list or a tuple of the same length as `len(self.sublayers)`
        and it will be used to set the state of all sublayers.
    """
    if isinstance(state, dict) and state == GET_STATE_FROM_CACHE:
      return
    if not self._sublayers:
      self._state = state
    else:
      # When having sublayers, self._state just marks which are cached,
      # the actual weights are stored by sublayers.
      self._state = []
      for s in state:
        if isinstance(s, dict) and s == GET_STATE_FROM_CACHE:
          self._state.append(s)
        else:
          self._state.append(None)
      # Set sublayer states.
      n_layers = len(self.sublayers)
      if len(state) != n_layers:
        raise ValueError(
            f'Number of state elements ({len(state)}) does not equal the '
            f'number of sublayers ({n_layers}) in: {str(self)}.')
      for sublayer, sublayer_state in zip(self.sublayers, state):
        sublayer.state = sublayer_state
Sets the state of this layer and its sublayers.

    Args:
      state: the state to set; if layer has sublayers, state should be
        either a list or a tuple of the same length as `len(self.sublayers)`
        and it will be used to set the state of all sublayers.
    
trax/layers/base.py