class_docstr: Elmo-style weighted average representation.
class_docstr_tok: ["Elmo-style","weighted","average","representation","."]
code: def forward(self, last_layer_feats, padding_mask, all_layer_feats):
      # last_layer_feats: [BxTxD]
      # padding_mask: [BxT]
      if not self.training:
          msg = (
              f"Number of layers in input features = {len(all_layer_feats)}."
              f" Expected {self.num_layers} layers."
          )
          assert len(all_layer_feats) == self.num_layers, msg
  
      # Stack up all layers and reshape to (num_layers, features)
      all_layer_feats_stacked = torch.stack(all_layer_feats, dim=0)
      num_layers, *original_feat_shape = all_layer_feats_stacked.shape
      all_layer_feats_stacked_flat = all_layer_feats_stacked.view(
          num_layers, -1
      )
  
      # Weighted average
      normalized_weights = F.softmax(self.weights, dim=-1)
      weighted_avg_features = (
          normalized_weights.unsqueeze(-1)
          * all_layer_feats_stacked_flat
      ).sum(dim=0)
      weighted_avg_features = weighted_avg_features.view(
          *original_feat_shape
      )
  
      # Mean Pooling on weighted average features.
      return super().forward(weighted_avg_features, padding_mask)
  
imports: import contextlib
  import logging
  from argparse import Namespace
  from dataclasses import dataclass, field
  from typing import Any, Optional
  import torch
  import torch.nn
  import torch.nn.functional
  from omegaconf import II, MISSING, open_dict
  from fairseq import checkpoint_utils, tasks, utils
  from fairseq.dataclass import ChoiceEnum, FairseqDataclass
  from fairseq.dataclass.utils import convert_namespace_to_omegaconf
  from fairseq.models import BaseFairseqModel, FairseqEncoder, register_model
  from fairseq.models.wav2vec.wav2vec2 import MASKING_DISTRIBUTION_CHOICES, Wav2Vec2Config
  from fairseq.models.wav2vec.wav2vec2_asr import Embedding, Linear, Wav2VecEncoder, Wav2Vec2AsrConfig
  from fairseq.tasks import FairseqTask
label_desc: Use image node_id 2 for calling the LayerWeightedMeanPooling obj's underlying member method code with example usage: obj.forward(last_layer_feats, padding_mask, all_layer_feats) and returns: super