def UniformlySeek(name=None, host_id=None, n_hosts=None, dataset_size=None):  # pylint: disable=invalid-name
  """Sets each host at (dataset_size/n_hosts)-th of the dataset."""
  if not dataset_size:
    dataset_size = 2 ** 18  # 512 * 512
    logging.error(
        'No dataset size given to Uniformly seek, assuming: %d', dataset_size)
  assert name
  host_id = jax.process_index() if host_id is None else host_id
  n_hosts = n_hosts or jax.host_count()
  each_host = int(dataset_size / n_hosts)
  def _f(generator):
    # Each host seeks to the appropriate point in the dataset.
    num_to_seek = int(host_id * each_host)
    start_time = time.time()
    logging.info('Dataset[%s] host_id[%d] is seeking to position[%d]',
                 name, host_id, num_to_seek)
    for _ in range(num_to_seek):
      next(generator)
    logging.info('Dataset[%s] host_id[%d] reached position[%d]. '
                 'Time taken [%s] seconds',
                 name, host_id, num_to_seek, time.time() - start_time)
    for example in generator:
      yield example
  return _f
Sets each host at (dataset_size/n_hosts)-th of the dataset.
trax/data/inputs.py