def input_fn(params, mode, get_dataset_fns, weights=None, mix_batches=True):
  """Wrapper for getting a dataset and batching it."""
  # Load datasets.
  datasets = [fn(params, mode) for fn in get_dataset_fns]

  # Weird bug where tensor2tensor doens't like non-static shapes on CPU,
  # but is ok when it's on TPU. Luckily we evaluate on TPU in our case.
  if mode != tf_estimator.ModeKeys.PREDICT:
    drop_remainder = True
  elif not params["use_tpu"]:
    tf.logging.warning("Evaluation intended for TPU! Will have to drop "
                       "remainder in batching. Results may be slightly off "
                       "(or use predict_batch_size=1)")
    drop_remainder = True
  else:
    drop_remainder = False

  # Normalize weights.
  if weights is not None:
    weights = np.array(weights).astype(np.float32)
    weights = weights / np.sum(weights)

  # Batch first, mix later.
  if not mix_batches:
    datasets = [
        dset.batch(params["batch_size"], drop_remainder=drop_remainder)
        for dset in datasets
    ]
    dataset = tf.data.experimental.sample_from_datasets(datasets, weights)

  # Mix first, batch later.
  else:
    dataset = tf.data.experimental.sample_from_datasets(datasets, weights)
    dataset = dataset.batch(params["batch_size"], drop_remainder=drop_remainder)

  dataset = dataset.prefetch(params["prefetch_batches"])

  return dataset
Wrapper for getting a dataset and batching it.
language/capwap/datasets/__init__.py