class_docstr: Represents an ALiBi attention mask as described in
  :cite:t:`https://doi.org/10.48550/arxiv.2108.12409`.
  
  *Shape:* :math:`(H,S,S_{kv})`, where :math:`H` is the number of attention
  heads, :math:`S` is the sequence length, and :math:`S_{kv}` is the key/value
  sequence length.
class_docstr_tok: ["Represents","an","ALiBi","attention","mask","as","described","in",":","cite",":","t",":","`","https",":","\/\/doi.org\/10.48550\/arxiv.2108.12409","`",".","*","Shape",":","*",":","math",":","`","(","H",",","S",",","S_","{","kv","}",")","`",",","where",":","math",":","`","H","`","is","the","number","of","attention","heads",",",":","math",":","`","S","`","is","the","sequence","length",",","and",":","math",":","`","S_","{","kv","}","`","is","the","key\/value","sequence","length","."]
code: def _do_materialize(self) -> Tensor:
      attn_len = (
          self.seq_len if self.attn_len is None else self.attn_len
      )
  
      # (H)
      powers = torch.arange(
          1, 1 + self.num_attn_heads, device=self.device
      )
  
      # (H)
      slopes = torch.pow(2 ** (-8 / self.num_attn_heads), powers)
  
      # (S_kv)
      steps = torch.arange(self.key_len, device=self.device)
  
      # (S_kv) -> (H, S, S_kv)
      steps = steps[None, None, :].expand(
          self.num_attn_heads, attn_len, -1
      )
  
      # (H, S, S_kv) * (H, 1, 1) -> (H, S, S_kv)
      mask = steps * slopes[:, None, None]
  
      mask = mask.to(self.dtype)
  
      # If the attention length is 1, avoid constructing the causal mask.
      if attn_len == 1:
          # Ensure that we do not attend to keys beyond the sequence length.
          if (causal := self.key_len - self.seq_len) > 0:
              mask[:, :, -causal:] = -torch.inf
      else:
          # (S, S_kv)
          causal_mask = _create_causal_attention_mask(
              self.seq_len,
              self.key_len,
              attn_len,
              None,
              self.device,
              self.dtype,
          )
  
          # (H, S, S_kv) + (S, S_kv) -> (H, S, S_kv)
          mask = mask + causal_mask
  
      return mask
  
imports: from abc import ABC, abstractmethod
  from typing import Optional, Protocol, final
  import torch
  from torch import Tensor
  from fairseq2.nn.incremental_state import IncrementalStateBag
  from fairseq2.typing import DataType, Device, finaloverride
label_desc: Use image node_id 2 for calling the ALiBiMask obj's underlying member method code with example usage: obj._do_materialize() and returns: mask