Method init_state

axlearn/common/decoding.py:300–311

Initializes prefix merger state. Args: tokens: The initial live sequences, of shape [batch_size, num_decodes, max_decode_len]. When prefilling the decoding cache, this consists of the decoding prefixes padded to `max_decode_len`. Returns:

(self, *, tokens: Tensor)

    """

    def init_state(self, *, tokens: Tensor) -> NestedTensor:
        """Initializes prefix merger state.

        Args:
            tokens: The initial live sequences, of shape [batch_size, num_decodes, max_decode_len].
                When prefilling the decoding cache, this consists of the decoding prefixes padded to
                `max_decode_len`.

        Returns:
            The initial state.
        """
        raise NotImplementedError(type(self))

    def compute(self, state: NestedTensor) -> Tensor:
        """Computes the merge matrix.