def _prepare_weights(self, weights):
    """Replicate or shard weights for the number of devices requested."""
    if base.N_WEIGHTS_SHARDS > 1:
      if base.N_WEIGHTS_SHARDS % self._n_devices != 0:
        raise ValueError(f'Number of shards ({base.N_WEIGHTS_SHARDS}) must '
                         f'be a multiple of n_devices ({self._n_devices}).')
      return base.shard(weights, base.N_WEIGHTS_SHARDS)
    else:
      return for_n_devices(weights, self._n_devices)
Replicate or shard weights for the number of devices requested.
trax/layers/acceleration.py