def _sample_proportionally(inputs, weights):
  """Sample an element from the inputs list proportionally to weights.

  Args:
    inputs: a list, we will return one element of this list.
    weights: a sequence of numbers of the same length as inputs; we will sample
      the k-th input with probability weights[k] / sum(weights).

  Returns:
    an element from inputs.
  """
  l = len(inputs)
  weights = np.array(weights)
  if l != len(weights):
    raise ValueError(f'Inputs and weights must have the same length, but do not'
                     f': {l} != {len(weights)}')
  norm_weights = weights / np.sum(weights)
  # TODO(pkozakowski): Currently this is O(n). It can be sped up to O(log n) by
  # storing CDF and binsearching on it.
  idx = np.random.choice(l, p=norm_weights)
  return inputs[int(idx)]
Sample an element from the inputs list proportionally to weights.

  Args:
    inputs: a list, we will return one element of this list.
    weights: a sequence of numbers of the same length as inputs; we will sample
      the k-th input with probability weights[k] / sum(weights).

  Returns:
    an element from inputs.
  
trax/rl/task.py