Method beam_search_decode

axlearn/common/decoder.py:845–859

See configured `decoding` implementation for details.

(
        self,
        *,
        input_batch: Nested[Tensor],
        max_sequence_length: int,
        num_decodes: int,
        **kwargs,
    )

            # recursively search through potentials.
            # They backtrace to some anchor time step after exploring for t steps.
            # This requires tracking time_step separately from the attention time_step.
            time_step=(cached_states["time_step"] + step_len),
        )
        if cfg.attention_mask is not None:
            updated_states["input_ids"] = updated_inputs
        return updated_states, outputs

    def beam_search_decode(
        self,
        *,
        input_batch: Nested[Tensor],
        max_sequence_length: int,
        num_decodes: int,
        **kwargs,
    ):
        """See configured `decoding` implementation for details."""
        return self._decoding.beam_search_decode(
            input_batch=input_batch,