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

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

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