def addition_input_stream(
    vocab_size=gin.REQUIRED, batch_size=gin.REQUIRED, min_length=gin.REQUIRED,
    max_length=gin.REQUIRED, pad_to_multiple=32, encdec=False):
  """Data stream for the add problem: <S>x+y<S>(x+y).

  Args:
    vocab_size: how many symbols to use.
    batch_size: how large are the batches.
    min_length: minimal length of w.
    max_length: maximal length of w.
    pad_to_multiple: int, pad length to be multiple of this number.
    encdec: bool, if True return encoder-decoder style inputs (default: False)

  Returns:
    python generator of tuples of data examples
  """
  base = vocab_size - 3  # We use 0 to pad, base+1 as "+" and base+2 as "<S>".
  def single_example(max_length, min_length):
    """Generate a stream of random mini-batches."""
    add_len = (min_length - 1) // 2
    l1 = np.random.randint((max_length - add_len + 1) // 2) + add_len
    l2 = np.random.randint(max_length - l1 - 1) + 1
    n1 = random_number_lower_endian(l1, base)
    n2 = random_number_lower_endian(l2, base)
    result = lower_endian_to_number(n1, base) + lower_endian_to_number(
        n2, base)
    inp = n1 + [base] + n2
    tgt = number_to_lower_endian(result, base)
    if encdec:
      x = [i + 1 for i in inp]
      y = [i + 1 for i in tgt]
      weights = [1] * len(tgt)
      candidate_example = (np.array(x), np.array(y), np.array(weights))
      if any(len(sample) > max_length for sample in candidate_example):
        # sample too long, try again
        return single_example(max_length, min_length)
      return (np.array(x), np.array(y), np.array(weights))
    else:
      x = [base+2] + [i+1 for i in inp] + [base+2] + [i+1 for i in tgt]
      weights = ([0] * (len(inp) + 2)) + ([1] * len(tgt))
      return (np.array(x), np.array(x), np.array(weights))

  def batches(max_length, min_length):
    """Batches of examples."""
    if max_length < 3 or min_length < 3:
      raise ValueError('Maximum/minimum length must be at least 3.')
    while True:
      ex = [single_example(max_length, min_length) for _ in range(batch_size)]
      padded_batch = [pad_to_max_dims(x, boundary=pad_to_multiple,
                                      strict_pad_on_len=True)
                      for x in zip(*ex)]
      yield tuple(padded_batch)

  return batches(max_length, min_length)
Data stream for the add problem: <S>x+y<S>(x+y).

  Args:
    vocab_size: how many symbols to use.
    batch_size: how large are the batches.
    min_length: minimal length of w.
    max_length: maximal length of w.
    pad_to_multiple: int, pad length to be multiple of this number.
    encdec: bool, if True return encoder-decoder style inputs (default: False)

  Returns:
    python generator of tuples of data examples
  
trax/data/inputs.py