class_docstr: Transformer-based, non-autoregressive model, operates on multiple streams of audio tokens (MAGNeT).
  Args:
      subcodes_context (int): The number of timesteps attended in the self-attention blocks of codebooks > 0.
                              When set to -1, attention is unrestricted and all timesteps are attended. Defaults to 5.
      compression_model_framerate (int): frame rate of the audio tokenizer.
      segment_duration (int): Sample length in seconds.
      span_len (int): Determines the length of masking spans. This is the minimal length of consecutive masked tokens,
                      for both training and inference. Defaults to 3.
      **kwargs: Additional parameters for the LMModel.
class_docstr_tok: ["Transformer-based",",","non-autoregressive","model",",","operates","on","multiple","streams","of","audio","tokens","(","MAGNeT",")",".","Args",":","subcodes_context","(","int",")",":","The","number","of","timesteps","attended","in","the","self-attention","blocks","of","codebooks",">","0",".","When","set","to","-1",",","attention","is","unrestricted","and","all","timesteps","are","attended",".","Defaults","to","5.","compression_model_framerate","(","int",")",":","frame","rate","of","the","audio","tokenizer",".","segment_duration","(","int",")",":","Sample","length","in","seconds",".","span_len","(","int",")",":","Determines","the","length","of","masking","spans",".","This","is","the","minimal","length","of","consecutive","masked","tokens",",","for","both","training","and","inference",".","Defaults","to","3",".","*","*","kwargs",":","Additional","parameters","for","the","LMModel","."]
docstr: Creates a restricted attention mask (local attention map) where the context
     is determined by self.subcodes_context.
  Args:
      seq_len (int): token sequence length.
      device (torch.device): device of the output tensor.
      dtype (torch.dtype): data type of the output tensor.
  Returns:
      torch.Tensor: The restricted attention mask.
docstr_tok: ["Creates","a","restricted","attention","mask","(","local","attention","map",")","where","the","context","is","determined","by","self.subcodes_context",".","Args",":","seq_len","(","int",")",":","token","sequence","length",".","device","(","torch.device",")",":","device","of","the","output","tensor",".","dtype","(","torch.dtype",")",":","data","type","of","the","output","tensor",".","Returns",":","torch.Tensor",":","The","restricted","attention","mask","."]
code: def restricted_context_attn_mask(
      self, seq_len: int, device: torch.device, dtype: torch.dtype
  ) -> torch.Tensor:
      """Creates a restricted attention mask (local attention map) where the context
         is determined by self.subcodes_context.
      Args:
          seq_len (int): token sequence length.
          device (torch.device): device of the output tensor.
          dtype (torch.dtype): data type of the output tensor.
      Returns:
          torch.Tensor: The restricted attention mask.
      """
      # Return a context restricted non-causal att mask
      queries_pos = torch.arange(seq_len, device=device).view(-1, 1)
      keys_pos = torch.arange(seq_len, device=device).view(1, -1)
  
      delta = queries_pos - keys_pos
      valid = torch.abs(delta) <= self.subcodes_context
      return torch.where(
          valid,
          torch.zeros([], device=device, dtype=dtype),
          torch.full([], float("-inf"), device=device, dtype=dtype),
      )
  
file_path: audiocraft/audiocraft/models/lm_magnet.py
imports: import logging
  import math
  import typing
  import torch
  import numpy
  from ..utils import utils
  from ..modules.conditioners import ClassifierFreeGuidanceDropout, ConditioningAttributes, ConditionType
  from .lm import LMModel
label_desc: Use image node_id 2 for calling the MagnetLMModel obj's underlying member method code with example usage: obj.restricted_context_attn_mask(seq_len, device, dtype) and returns: torch