Method forward

axlearn/common/decoder.py:644–691

Computes decoder hidden states from input ids and cross attention hidden states. Args: input_batch: A dict containing: * input_ids: An int Tensor of shape [batch_size, target_len]. Values should be in the range [0, vocab_size).

(
        self,
        input_batch: Nested[Tensor],
        *,
        cross_attention_data: Optional[Tensor] = None,
        cross_attention_logit_biases: Optional[Tensor] = None,
        **kwargs,
    )

        return logits

    def forward(
        self,
        input_batch: Nested[Tensor],
        *,
        cross_attention_data: Optional[Tensor] = None,
        cross_attention_logit_biases: Optional[Tensor] = None,
        **kwargs,
    ) -> dict[str, Tensor]:
        """Computes decoder hidden states from input ids and cross attention hidden states.

        Args:
            input_batch: A dict containing:
                * input_ids: An int Tensor of shape [batch_size, target_len].
                    Values should be in the range [0, vocab_size).
                * input_segment_ids: An optional Tensor of same shape as `input_ids` with values in
                    [0, num_segments]. Tokens are only allowed to attend to other tokens within the
                    same segment. input_segment_ids == 0 represents paddings. If None, inferred from
                    input_ids != pad_token_id.
                * token_type_ids: An optional int Tensor of shape [batch_size, target_len].
                    Values should be in the range [0, type_vocab_size).
                * positions: An optional int Tensor of shape [batch_size, target_len].
                    If None, assumed to be jnp.arange(target_len) for each sequence.
            cross_attention_data: A float Tensor of shape [batch_size, source_len, hidden_dim].
            cross_attention_logit_biases: A Tensor of shape [batch_size, target_len, source_len].
                A -inf represents a disconnected position pair.

        Returns:
            A dict containing:
                hidden_states: A float Tensor of shape [batch_size, target_len, hidden_dim].
        """
        validate_contains_paths(input_batch, paths=["input_ids"])
        input_ids = input_batch["input_ids"]
        input_segment_ids = input_batch.get("input_segment_ids", None)
        positions = input_batch.get("positions", None)

        _, output = self._forward_for_mode(
            mode=ForwardMode.FORWARD,
            input_batch=input_batch,
            # [batch_size, num_heads, seq_len, seq_len].
            self_attention_logit_biases=self.compute_attention_logit_biases(
                input_ids, segment_ids=input_segment_ids, positions=positions
            ),
            cross_attention_data=cross_attention_data,
            cross_attention_logit_biases=cross_attention_logit_biases,
            cached_states=None,
            **kwargs,
        )
        return output

    @nowrap
    def init_states(