def replicate_state(self, state):
    """Sets the state of the sublayer and replicates it for this layer."""
    self.sublayer.state = state
    self._state = for_n_devices(state, self._n_devices)
Sets the state of the sublayer and replicates it for this layer.
trax/layers/acceleration.py