message: [tf.data] Optimize `tf.contrib.data.sample_from_datasets()` when the weights are not a dataset.
  Previously, we were recomputing the logits from the weights for each element, which is only necessary when the weights can differ for each element.
diff: @@ -216,25 +216,46 @@ def sample_from_datasets(datasets, weights=None, seed=None):
  length of the `datasets` element.
  """
  num_datasets = len(datasets)
  + if not isinstance(weights, dataset_ops.Dataset):
  if weights is None:
  - weights = dataset_ops.Dataset.from_tensors([1.0] * num_datasets).repeat()
  - elif not isinstance(weights, dataset_ops.Dataset):
  + # Select inputs with uniform probability.
  + logits = [[1.0] * num_datasets]
  + else:
  + # Use the given `weights` as the probability of choosing the respective
  + # input.
  weights = ops.convert_to_tensor(weights, name="weights")
  if weights.dtype not in (dtypes.float32, dtypes.float64):
  raise TypeError("`weights` must be convertible to a tensor of "
  "`tf.float32` or `tf.float64` elements.")
  if not weights.shape.is_compatible_with([num_datasets]):
  - raise ValueError("`weights` must be a vector of length `len(datasets)`.")
  - weights = dataset_ops.Dataset.from_tensors(weights).repeat()
  + raise ValueError(
  + "`weights` must be a vector of length `len(datasets)`.")
  +
  + # The `stateless_multinomial()` op expects log-probabilities, as opposed
  + # to weights.
  + logits = array_ops.expand_dims(math_ops.log(weights, name="logits"), 0)
  +
  + def select_dataset_constant_logits(seed):
  + return array_ops.squeeze(
  + stateless.stateless_multinomial(logits, 1, seed=seed), axis=[0, 1])
  +
  + selector_input = random_ops.RandomDataset(seed).batch(2).map(
  + select_dataset_constant_logits)
  + else:
  + # Use each element of the given `weights` dataset as the probability of
  + # choosing the respective input.
  # The `stateless_multinomial()` op expects log-probabilities, as opposed to
  # weights.
  logits_ds = weights.map(lambda *p: math_ops.log(p, name="logits"))
  - def select_dataset(logits, seed):
  +
  + def select_dataset_varying_logits(logits, seed):
  return array_ops.squeeze(
  stateless.stateless_multinomial(logits, 1, seed=seed), axis=[0, 1])
  +
  selector_input = dataset_ops.Dataset.zip(
  - (logits_ds, random_ops.RandomDataset(seed).batch(2))).map(select_dataset)
  + (logits_ds, random_ops.RandomDataset(seed).batch(2)
  + )).map(select_dataset_varying_logits)
  return _DirectedInterleaveDataset(selector_input, datasets)
  