Method _pad

axlearn/common/decoder.py:461–473

Accept token IDs input tensor and pad if necessary to max_sequence_length.

(prefix: Tensor, *, max_sequence_length: int, pad_id: int)

                    dtype=prefix.dtype,
                ),
            ],
            axis=1,
        )


# TODO(gyin): Add unittest for Decoder forward
class Decoder(BaseLayer):
    """Construct a decoder transformer to output hidden states and logits based on lm head."""

    @config_class
    class Config(BaseLayer.Config):
        """Configures Decoder."""

        # DEPRECATED, because `attention_mask` uses quadratic memory, even with Flash Attention.
        # Please use `attention.mask`, which constructs the mask procedurally.
        attention_mask: Optional[AttentionLogitBiasLayer.Config] = None