def _very_simple_data(output_dim=1, input_dim=1):
  """"Returns stream of labeled data that maps small integers to constant pi."""
  inputs_batch = np.arange(8).reshape((8, 1))  # 8 items per batch
  inputs_batch = np.concatenate([inputs_batch] * input_dim, axis=1)
  targets_batch = np.pi * np.ones((8, output_dim))
  labeled_batch = (inputs_batch, targets_batch, np.ones_like(targets_batch))
  while True:
    yield labeled_batch
"Returns stream of labeled data that maps small integers to constant pi.
trax/supervised/training_test.py