class_docstr: This class implements multi-head HyperMixing.
  It is an implementation of the token-mixing component in HyperMixer, a linear
  time drop-in replacement for self-attention. In contrast to the original HyperMixer,
  this module supports multiple heads, which improves the expressiveness of the model
  while decreasing the number of parameters.
  
  Reference: https://arxiv.org/abs/2203.03691
  
  Arguments
  ----------
  input_output_dim : int
      number of features in keys, queries, and values
  hypernet_size : int
      determines the size of the hidden layer of the token-mixing MLP.
  tied : bool
      If True, then the generated weight matrices of the token-mixing MLP are tied.
  num_heads : int
      parallel token-mixing MLPs.
  fix_tm_hidden_size : bool
      If True, the hidden-layer size is equal to hypernet_size rather than hypernet_size / num_heads.
  max_length : int
      Maximum number of input tokens. Needed for generating sufficiently large position embeddings.
  
  Example
  -------
  >>> import torch
  >>> inputs = torch.rand([8, 60, 512])
  >>> net = HyperMixing(512, 2048, num_heads=8)
  >>> outputs, attn = net(inputs, inputs, inputs)
  >>> outputs.shape
  torch.Size([8, 60, 512])
class_docstr_tok: ["This","class","implements","multi-head","HyperMixing",".","It","is","an","implementation","of","the","token-mixing","component","in","HyperMixer",",","a","linear","time","drop-in","replacement","for","self-attention",".","In","contrast","to","the","original","HyperMixer",",","this","module","supports","multiple","heads",",","which","improves","the","expressiveness","of","the","model","while","decreasing","the","number","of","parameters",".","Reference",":","https",":","\/\/arxiv.org\/abs\/2203.03691","Arguments","--","--","--","--","--","input_output_dim",":","int","number","of","features","in","keys",",","queries",",","and","values","hypernet_size",":","int","determines","the","size","of","the","hidden","layer","of","the","token-mixing","MLP",".","tied",":","bool","If","True",",","then","the","generated","weight","matrices","of","the","token-mixing","MLP","are","tied",".","num_heads",":","int","parallel","token-mixing","MLPs",".","fix_tm_hidden_size",":","bool","If","True",",","the","hidden-layer","size","is","equal","to","hypernet_size","rather","than","hypernet_size","\/","num_heads",".","max_length",":","int","Maximum","number","of","input","tokens",".","Needed","for","generating","sufficiently","large","position","embeddings",".","Example","--","--","--","-",">",">",">","import","torch",">",">",">","inputs","=","torch.rand","(","[","8",",","60",",","512","]",")",">",">",">","net","=","HyperMixing","(","512",",","2048",",","num_heads=8",")",">",">",">","outputs",",","attn","=","net","(","inputs",",","inputs",",","inputs",")",">",">",">","outputs.shape","torch.Size","(","[","8",",","60",",","512","]",")"]
docstr: The signature of this method is deliberately chosen to be the same as for
  sb.nnet.attention.MultiHeadAttention for compatibility within SpeechBrain.
  
  NOTE: key, value, attn_mask and pos_embs have no effect. Query is used for
  all three. Thus, the module should only be used to replace self-attention at the moment.
  
  Arguments
  ----------
  query : torch.Tensor
      (B, L, E) where L is the target sequence length,
      B is the batch size, E is the embedding dimension.
  key : torch.Tensor
      (B, S, E) where S is the source sequence length,
      B is the batch size, E is the embedding dimension.
      Currently unused. All
  value : torch.Tensor
      (B, S, E) where S is the source sequence length,
      B is the batch size, E is the embedding dimension.
      Currently unused.
  attn_mask : torch.Tensor, optional
      NOTE: Currently has NO effect.
  key_padding_mask : torch.Tensor, optional
      (B, S) where B is the batch size, S is the source sequence
      length. If a ByteTensor is provided, the non-zero positions will
      be ignored while the position with the zero positions will be
      unchanged. If a BoolTensor is provided, the positions with the
      value of True will be ignored while the position with the value
      of False will be unchanged.
  return_attn_weights: torch.Tensor, optional
      NOTE: Currently has NO effect.
  pos_embs: torch.Tensor, optional
      NOTE: Currently has NO effect.
  
  Outputs
  -------
  attn_output : torch.Tensor
      (B, L, E) where L is the target sequence length, B is the
      batch size, E is the embedding dimension.
  attn_output_weights : torch.Tensor
      (B, L, S) where B is the batch size, L is the target
      sequence length, S is the source sequence length.
      NOTE: always returns all zeros.
docstr_tok: ["The","signature","of","this","method","is","deliberately","chosen","to","be","the","same","as","for","sb.nnet.attention.MultiHeadAttention","for","compatibility","within","SpeechBrain",".","NOTE",":","key",",","value",",","attn_mask","and","pos_embs","have","no","effect",".","Query","is","used","for","all","three",".","Thus",",","the","module","should","only","be","used","to","replace","self-attention","at","the","moment",".","Arguments","--","--","--","--","--","query",":","torch.Tensor","(","B",",","L",",","E",")","where","L","is","the","target","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","key",":","torch.Tensor","(","B",",","S",",","E",")","where","S","is","the","source","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","Currently","unused",".","All","value",":","torch.Tensor","(","B",",","S",",","E",")","where","S","is","the","source","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","Currently","unused",".","attn_mask",":","torch.Tensor",",","optional","NOTE",":","Currently","has","NO","effect",".","key_padding_mask",":","torch.Tensor",",","optional","(","B",",","S",")","where","B","is","the","batch","size",",","S","is","the","source","sequence","length",".","If","a","ByteTensor","is","provided",",","the","non-zero","positions","will","be","ignored","while","the","position","with","the","zero","positions","will","be","unchanged",".","If","a","BoolTensor","is","provided",",","the","positions","with","the","value","of","True","will","be","ignored","while","the","position","with","the","value","of","False","will","be","unchanged",".","return_attn_weights",":","torch.Tensor",",","optional","NOTE",":","Currently","has","NO","effect",".","pos_embs",":","torch.Tensor",",","optional","NOTE",":","Currently","has","NO","effect",".","Outputs","--","--","--","-","attn_output",":","torch.Tensor","(","B",",","L",",","E",")","where","L","is","the","target","sequence","length",",","B","is","the","batch","size",",","E","is","the","embedding","dimension",".","attn_output_weights",":","torch.Tensor","(","B",",","L",",","S",")","where","B","is","the","batch","size",",","L","is","the","target","sequence","length",",","S","is","the","source","sequence","length",".","NOTE",":","always","returns","all","zeros","."]
code: def forward(
      self,
      query,
      key,
      value,
      attn_mask: Optional[torch.Tensor] = None,
      key_padding_mask: Optional[torch.Tensor] = None,
      return_attn_weights: Optional[bool] = True,
      pos_embs: Optional[torch.Tensor] = None,
  ):
      """
      The signature of this method is deliberately chosen to be the same as for
      sb.nnet.attention.MultiHeadAttention for compatibility within SpeechBrain.
  
      NOTE: key, value, attn_mask and pos_embs have no effect. Query is used for
      all three. Thus, the module should only be used to replace self-attention at the moment.
  
      Arguments
      ----------
      query : torch.Tensor
          (B, L, E) where L is the target sequence length,
          B is the batch size, E is the embedding dimension.
      key : torch.Tensor
          (B, S, E) where S is the source sequence length,
          B is the batch size, E is the embedding dimension.
          Currently unused. All
      value : torch.Tensor
          (B, S, E) where S is the source sequence length,
          B is the batch size, E is the embedding dimension.
          Currently unused.
      attn_mask : torch.Tensor, optional
          NOTE: Currently has NO effect.
      key_padding_mask : torch.Tensor, optional
          (B, S) where B is the batch size, S is the source sequence
          length. If a ByteTensor is provided, the non-zero positions will
          be ignored while the position with the zero positions will be
          unchanged. If a BoolTensor is provided, the positions with the
          value of True will be ignored while the position with the value
          of False will be unchanged.
      return_attn_weights: torch.Tensor, optional
          NOTE: Currently has NO effect.
      pos_embs: torch.Tensor, optional
          NOTE: Currently has NO effect.
  
      Outputs
      -------
      attn_output : torch.Tensor
          (B, L, E) where L is the target sequence length, B is the
          batch size, E is the embedding dimension.
      attn_output_weights : torch.Tensor
          (B, L, S) where B is the batch size, L is the target
          sequence length, S is the source sequence length.
          NOTE: always returns all zeros.
      """
  
      # NOTE: We are ignoring keys and values, because HyperMixing can only be used in the encoder atm (where it's all the same)
      out = query
  
      bsize = out.size(0)
      seq_len = out.size(1)
  
      if key_padding_mask is not None:
          float_mask = (
              torch.logical_not(key_padding_mask).unsqueeze(-1).float()
          )
          out = out * float_mask
  
      # add position embedding before passing to hypernetwork
      hyp_input = out + self.positional_encoding(out)
      W1, W2 = self.hyper(
          hyp_input
      )  # [bsize, num_heads, seq_len, hypernet_size // num_heads]
  
      if key_padding_mask is not None:
          # mask the weights
          W1 = W1 * float_mask.unsqueeze(1)
          W2 = W2 * float_mask.unsqueeze(1)
  
      # reshape the num_heads into the batch dimension for parallelizing
      out = out.transpose(1, 2)  # [bsize, input_output_dim, seq_len]
      out = out.reshape(
          (
              bsize * self.num_heads,
              self.input_output_dim // self.num_heads,
              seq_len,
          )
      )  # [bsize * num_heads, input_output_dim // num_heads, seq_len]
      W1 = W1.reshape((bsize * self.num_heads, seq_len, -1))
      W2 = W2.reshape((bsize * self.num_heads, seq_len, -1))
  
      # we stick the token-mixing MLP together manually
      out = self._mlp_pass_from_components(out, W1, W2, self.activation)
  
      # concatenate heads
      out = out.reshape((bsize, self.input_output_dim, seq_len))
  
      # transpose back
      out = out.transpose(1, 2)
  
      # apply layer norm on outputs of the TM-MLP
      out = self.layer_norm(out)
  
      dummy_att_weights = torch.zeros(
          (bsize, seq_len, seq_len), device=out.device
      )
      return out, dummy_att_weights
  
file_path: speechbrain/speechbrain/lobes/models/transformer/hypermixing.py
imports: from typing import Optional
  import math
  import torch
  from torch import nn
label_desc: Use image node_id 3 for calling the HyperMixing obj's underlying member method code with example usage: obj.forward(query, key, value, attn_mask, key_padding_mask, return_attn_weights, pos_embs) and returns: out,  dummy_att_weights