Merge pull request #14240 from sferik/flat_map

Replace map.flatten(1) with flat_map