Function _beam_init

axlearn/common/decoding.py:373–429

Initializes the beam search state data structure. Args: inputs: An int tensor of shape [batch_size, length] where length <= max_decode_length. time_step: Initial time steps for decoding of shape [batch_size]. batch_size: Size of batch. beam_size: Number of hypoth

(
    *,
    inputs: Tensor,
    time_step: Tensor,
    beam_size: int,
    max_decode_len: int,
    cache: NestedTensor,
    pad_id: int,
    prefix_merger: Optional[PrefixMerger] = None,
)

def _beam_init(
    *,
    inputs: Tensor,
    time_step: Tensor,
    beam_size: int,
    max_decode_len: int,
    cache: NestedTensor,
    pad_id: int,
    prefix_merger: Optional[PrefixMerger] = None,
) -> _BeamState:
    """Initializes the beam search state data structure.

    Args:
        inputs: An int tensor of shape [batch_size, length] where length <= max_decode_length.
        time_step: Initial time steps for decoding of shape [batch_size].
        batch_size: Size of batch.
        beam_size: Number of hypotheses per beam (aka num_decodes).
        max_decode_len: The maximum length of the sequence to be generated.
        cache: State of the decoder model.
        pad_id: Token ID associated with padded input.
        prefix_merger: Optional prefix merger.

    Returns:
        Fully initialized _BeamState.

    Raises:
        ValueError: If inputs has an invalid shape.
    """
    if inputs.shape[0] != time_step.shape[0]:
        raise ValueError(
            f"Expected inputs.shape[0] ({inputs.shape[0]}) "
            f"== time_step.shape[0] ({time_step.shape[0]})."
        )
    if inputs.shape[1] > max_decode_len:
        raise ValueError(
            f"Expected inputs.shape[1] ({inputs.shape[1]}) <= max_decode_len ({max_decode_len})."
        )
    batch_size = inputs.shape[0]
    live_seqs = jnp.full((batch_size, beam_size, max_decode_len), pad_id, dtype=jnp.int32)
    # Inputs are the prefix we will use for teacher forcing.
    live_seqs = live_seqs.at[:, :, : inputs.shape[1]].set(inputs[:, None, :])

    prefix_merger_state = None
    if prefix_merger is not None:
        prefix_merger_state = prefix_merger.init_state(tokens=live_seqs)

    return _BeamState(
        cur_index=time_step,
        # Handle first time step by masking out scores of all but the top hypothesis in the beam.
        live_scores=jnp.tile(jnp.array([0.0] + [NEG_INF] * (beam_size - 1)), [batch_size, 1]),
        finished_scores=jnp.ones((batch_size, beam_size)) * NEG_INF,
        live_seqs=live_seqs,
        finished_seqs=jnp.zeros((batch_size, beam_size, max_decode_len), jnp.int32),
        # Expand cache to num_decodes size.
        cache=vectorized_tree_map(lambda x: add_decoding_dim(x, beam_size), cache),
        prefix_merger=prefix_merger_state,
    )