class_docstr: ## Self-Attention Layer $   ext{A\small{TTN}}$
  
  This applies causal and non-causal [multi-headed self-attention](../mha.html).
class_docstr_tok: ["#","#","Self-Attention","Layer","$","ext","{","A\\small","{","TTN","}","}","$","This","applies","causal","and","non-causal","[","multi-headed","self-attention","]","(","..","\/mha.html",")","."]
docstr: ### Mask the attention layer for causal attention
  
  * `attn` is the attention matrix of shape `[batch_size, n_heads, seq_len, seq_len]`
docstr_tok: ["#","#","#","Mask","the","attention","layer","for","causal","attention","*","`","attn","`","is","the","attention","matrix","of","shape","`","[","batch_size",",","n_heads",",","seq_len",",","seq_len","]","`"]
code: def mask_attention(self, attn: torch.Tensor):
      """
      ### Mask the attention layer for causal attention
  
      * `attn` is the attention matrix of shape `[batch_size, n_heads, seq_len, seq_len]`
      """
  
      # No masking for non-causal attention
      if not self.is_causal:
          return attn
  
      # Create a triangular mask
      mask = torch.tril(attn.new_ones(attn.shape[-2:]))
      # Filter by the mask
      return attn.masked_fill(mask == 0, float("-inf"))
  
imports: import math
  from typing import Set
  import torch
  from torch import nn
  from labml.logger import inspect
label_desc: Use image node_id 2 for calling the SelfAttention obj's underlying member method code with example usage: obj.mask_attention(attn) and returns: attn, attn