def embed_documents_using_multiple_tpu_workers(
    shard_paths, shard_sizes, hub_module_spec,
    featurizer, tpu_workers,
    batch_size, num_tpu_cores_per_worker):
  """Embeds documents using multiple TPU workers.

  Args:
    shard_paths: a list of file paths, each specifying a GZIP'd TFRecord file
      containing documents stored as TF Examples. Doc embeddings will be
      concatenated in the order of shard_paths.
    shard_sizes: a list parallel to shard_paths, specifying the number of
      documents in each shard.
    hub_module_spec: path to the Hub module that will be used to embed the
      documents.
    featurizer: a Featurizer used to convert documents into Tensor features.
    tpu_workers: list of addresses of available TPU workers.
    batch_size: each TPU worker embeds documents in batches of this size.
    num_tpu_cores_per_worker: number of cores to use on each TPU worker.

  Returns:
    a [total_docs, embed_dim] Numpy array.
  """
  num_shards = len(shard_paths)
  num_tpu_workers = len(tpu_workers)

  tpu_config = tf_estimator.tpu.TPUConfig(
      iterations_per_loop=1,  # This seems to be ignored by predict().
      num_shards=num_tpu_cores_per_worker)

  # Distribute the data shards as evenly as possible among the workers.
  num_shards_per_worker = [num_shards // num_tpu_workers] * num_tpu_workers
  for worker_idx in range(num_shards % num_tpu_workers):
    num_shards_per_worker[worker_idx] += 1

  worker_kwargs = []
  shards_assigned = 0
  for k, num_shards_k in enumerate(num_shards_per_worker):
    worker_kwargs.append({
        'tpu_run_config':
            tf_estimator.tpu.RunConfig(
                master=tpu_workers[k], tpu_config=tpu_config),
        'shard_paths':
            shard_paths[shards_assigned:shards_assigned + num_shards_k],
        'shard_sizes':
            shard_sizes[shards_assigned:shards_assigned + num_shards_k],
        'hub_module_spec': hub_module_spec,
        'featurizer': featurizer,
        'batch_size': batch_size,
    })
    shards_assigned += num_shards_k

  # All shards should be assigned.
  assert shards_assigned == num_shards

  # Run all workers in parallel via separate threads.
  with futures.ThreadPoolExecutor(max_workers=num_tpu_workers) as executor:
    # A list of [num_docs_per_worker, embed_dim] Numpy arrays.
    embeds_list = list(
        executor.map(lambda kwargs: embed_documents(**kwargs), worker_kwargs))

    # A [total_docs, embed_dim] Numpy array.
    embeds = np.concatenate(embeds_list, axis=0)

  return embeds
Embeds documents using multiple TPU workers.

  Args:
    shard_paths: a list of file paths, each specifying a GZIP'd TFRecord file
      containing documents stored as TF Examples. Doc embeddings will be
      concatenated in the order of shard_paths.
    shard_sizes: a list parallel to shard_paths, specifying the number of
      documents in each shard.
    hub_module_spec: path to the Hub module that will be used to embed the
      documents.
    featurizer: a Featurizer used to convert documents into Tensor features.
    tpu_workers: list of addresses of available TPU workers.
    batch_size: each TPU worker embeds documents in batches of this size.
    num_tpu_cores_per_worker: number of cores to use on each TPU worker.

  Returns:
    a [total_docs, embed_dim] Numpy array.
  
language/realm/retrieval.py