def test_truncate_to_length_len_map(self, len_map, out_shapes):
    """Tests that truncation occurs when len_map values are small enough."""
    def data_stream():
      while True:
        yield (np.zeros((1, 5)), np.ones((1, 5)))
    stream_fn = data.inputs.TruncateToLength(len_map=len_map)
    y0, y1 = next(stream_fn(data_stream()))
    self.assertEqual(y0.shape, out_shapes[0])
    self.assertEqual(y1.shape, out_shapes[1])
Tests that truncation occurs when len_map values are small enough.
trax/data/inputs_test.py