def replicate_weights(self, weights):
    """Sets the weights of the sublayer and replicates them for this layer."""
    self.sublayer.weights = weights
    self._weights = self._prepare_weights(weights)
Sets the weights of the sublayer and replicates them for this layer.
trax/layers/acceleration.py