def _test_dataset_ints(inp_lengths, tgt_lengths):
  """Create a test dataset of int64 tensors of given shapes."""

  def generator():
    for inp_len, tgt_len in zip(inp_lengths, tgt_lengths):
      inp = np.ones([inp_len], dtype=np.int64)
      tgt = np.ones([tgt_len], dtype=np.int64)
      yield {'inputs': inp, 'targets': tgt}

  types = {'inputs': tf.int64, 'targets': tf.int64}
  shapes = {'inputs': tf.TensorShape([None]), 'targets': tf.TensorShape([None])}
  return tf.data.Dataset.from_generator(
      generator, output_types=types, output_shapes=shapes)
Create a test dataset of int64 tensors of given shapes.
trax/data/tf_inputs_test.py