class_docstr: Transformer-based language model on multiple streams of codes.
  
  Args:
      pattern_provider (CodebooksPatternProvider): Pattern provider for codebook interleaving.
      condition_provider (MusicConditioningProvider): Conditioning provider from metadata.
      fuser (ConditionFuser): Fuser handling the fusing of conditions with language model input.
      n_q (int): Number of parallel streams to model.
      card (int): Cardinality, vocabulary size.
      dim (int): Dimension of the transformer encoder.
      num_heads (int): Number of heads for the transformer encoder.
      hidden_scale (int): Scale for hidden feed forward dimension of the transformer encoder.
      norm (str): Normalization method.
      norm_first (bool): Use pre-norm instead of post-norm.
      emb_lr (float, optional): Embedding-specific learning rate.
      bias_proj (bool): Use bias for output projections.
      weight_init (str, optional): Method for weight initialization.
      depthwise_init (str, optional): Method for depthwise weight initialization.
      zero_bias_init (bool): If true and bias in Linears, initialize bias to zeros.
      cfg_dropout (float): Classifier-free guidance dropout.
      cfg_coef (float): Classifier-free guidance coefficient.
      attribute_dropout (dict): Attribute dropout probabilities.
      two_step_cfg (bool): Whether to run classifier free-guidance with 2 distinct steps.
      **kwargs: Additional parameters for the transformer encoder.
docstr: Apply language model on sequence and conditions.
  Given a tensor of sequence of shape [B, K, S] with K the number of codebooks and
  S the sequence steps, return the logits with shape [B, card, K, S].
  
  Args:
      indices (torch.Tensor): Indices of the codes to model.
      conditions (list of ConditioningAttributes): Conditions to use when modeling
          the given codes. Note that when evaluating multiple time with the same conditioning
          you should pre-compute those and pass them as `condition_tensors`.
      condition_tensors (dict[str, ConditionType], optional): Pre-computed conditioning
          tensors, see `conditions`.
      stage (int): The codebook level that is being predicted. Relevant for MAGNeT
          in which prediction is done in a codebook-by-codebook manner.
          Takes values in range(n_q), and ignored by default.
  Returns:
      torch.Tensor: Logits.
docstr_tok: ["Apply","language","model","on","sequence","and","conditions",".","Given","a","tensor","of","sequence","of","shape","[","B",",","K",",","S","]","with","K","the","number","of","codebooks","and","S","the","sequence","steps",",","return","the","logits","with","shape","[","B",",","card",",","K",",","S","]",".","Args",":","indices","(","torch.Tensor",")",":","Indices","of","the","codes","to","model",".","conditions","(","list","of","ConditioningAttributes",")",":","Conditions","to","use","when","modeling","the","given","codes",".","Note","that","when","evaluating","multiple","time","with","the","same","conditioning","you","should","pre-compute","those","and","pass","them","as","`","condition_tensors","`",".","condition_tensors","(","dict","[","str",",","ConditionType","]",",","optional",")",":","Pre-computed","conditioning","tensors",",","see","`","conditions","`",".","stage","(","int",")",":","The","codebook","level","that","is","being","predicted",".","Relevant","for","MAGNeT","in","which","prediction","is","done","in","a","codebook-by-codebook","manner",".","Takes","values","in","range","(","n_q",")",",","and","ignored","by","default",".","Returns",":","torch.Tensor",":","Logits","."]
code: def forward(
      self,
      sequence: torch.Tensor,
      conditions: tp.List[ConditioningAttributes],
      condition_tensors: tp.Optional[ConditionTensors] = None,
      stage: int = -1,
  ) -> torch.Tensor:
      """Apply language model on sequence and conditions.
      Given a tensor of sequence of shape [B, K, S] with K the number of codebooks and
      S the sequence steps, return the logits with shape [B, card, K, S].
  
      Args:
          indices (torch.Tensor): Indices of the codes to model.
          conditions (list of ConditioningAttributes): Conditions to use when modeling
              the given codes. Note that when evaluating multiple time with the same conditioning
              you should pre-compute those and pass them as `condition_tensors`.
          condition_tensors (dict[str, ConditionType], optional): Pre-computed conditioning
              tensors, see `conditions`.
          stage (int): The codebook level that is being predicted. Relevant for MAGNeT
              in which prediction is done in a codebook-by-codebook manner.
              Takes values in range(n_q), and ignored by default.
      Returns:
          torch.Tensor: Logits.
      """
      B, K, S = sequence.shape
      assert (
          K == self.num_codebooks
      ), "Sequence shape must match the specified number of codebooks"
      input_ = sum([self.emb[k](sequence[:, k]) for k in range(K)])
      if condition_tensors is None:
          assert (
              not self._is_streaming
          ), "Conditions tensors should be precomputed when streaming."
          # apply dropout modules
          conditions = self.cfg_dropout(conditions)
          conditions = self.att_dropout(conditions)
          tokenized = self.condition_provider.tokenize(conditions)
          # encode conditions and fuse, both have a streaming cache to not recompute when generating.
          condition_tensors = self.condition_provider(tokenized)
      else:
          assert (
              not conditions
          ), "Shouldn't pass both conditions and condition_tensors."
  
      input_, cross_attention_input = self.fuser(
          input_, condition_tensors
      )
  
      out = self.transformer(
          input_,
          cross_attention_src=cross_attention_input,
          src_mask=(
              self.attn_mask_per_stage[stage] if stage >= 0 else None
          ),
      )
      if self.out_norm:
          out = self.out_norm(out)
      logits = torch.stack(
          [self.linears[k](out) for k in range(K)], dim=1
      )  # [B, K, S, card]
  
      # remove the prefix from the model outputs
      if len(self.fuser.fuse2cond["prepend"]) > 0:
          logits = logits[:, :, -S:]
  
      return logits
  
imports: from dataclasses import dataclass
  from functools import partial
  import logging
  import math
  import typing
  import torch
  from torch import nn
  from ..utils import utils
  from ..modules.streaming import StreamingModule, State
  from ..modules.transformer import StreamingTransformer, create_norm_fn
  from ..modules.conditioners import ConditionFuser, ClassifierFreeGuidanceDropout, AttributeDropout, ConditioningProvider, ConditioningAttributes, ConditionType
  from ..modules.codebooks_patterns import CodebooksPatternProvider
  from ..modules.activations import get_activation_fn
label_desc: Use image node_id 5 for calling the LMModel obj's underlying member method code with example usage: obj.forward(sequence, conditions, condition_tensors, stage) and returns: logits