def init(self, input_signature):
    """Calls ``self.sublayer.init`` and replicates its values onto devices."""
    weights, state = self.sublayer.init(input_signature, use_cache=True)
    self._weights = self._prepare_weights(weights)
    self._state = for_n_devices(state, self._n_devices)
    return (self.weights, self.state)
Calls ``self.sublayer.init`` and replicates its values onto devices.
trax/layers/acceleration.py