pick_top_k function

Build an optree pipeline that selects up to the top k rows from each group in the given order.