def gather_best_candidate(nested, selected_idx_stacked, batch_axis=0):
            """Gathers the slices indexed by selected_idx_stacked from a potentially nested structure of tensors."""

            def gather_fn(tensor):
                gathered_tensor = tf.gather(params=tensor, indices=selected_idx_stacked, axis=batch_axis)
                return gathered_tensor

            return tf.nest.map_structure(gather_fn, nested)
Gathers the slices indexed by selected_idx_stacked from a potentially nested structure of tensors.
dpkd/transformers/src/transformers/generation/tf_utils.py