def random_minibatches(length_list):
    """Generate a stream of random mini-batches."""
    while True:
      length = random.choice(length_list)
      x = np.random.randint(low=1, high=vocab_size-1,
                            size=(batch_size, length))
      loss_weights = np.ones((batch_size, length))
      x = _pad_to_multiple_of(x, pad_to_multiple, 1)
      loss_weights = _pad_to_multiple_of(loss_weights, pad_to_multiple, 1)
      yield (x, x, loss_weights)  # Here inputs and targets are the same.
Generate a stream of random mini-batches.
trax/data/inputs.py