#5673 lists an extensive feature matrix for DeviceSegmentedTopK. There are various parameters that can either be segment-specific (e.g., for segment 0 I want the top-5 items, for segment 1 I want the top-2 items) or uniform across all the segments (e.g., I want to get the top-5 items for all my segments). We want to have the flexibility to cover various combinations of per-segment and uniform parameters across (e.g., segment_size, k, direction).