class_docstr: Latent variable modeling of attention head selection
class_docstr_tok: ["Latent","variable","modeling","of","attention","head","selection"]
code: def head_select(self, task_ids=None):
      # gumbel_sample
      self.head_samples = self.gumbel_sample(
          self.head_logits, tau=self.temp
      )
      # head select
      if self.select_strategy == "subset":
          self.subset_heads, self.subset_weights = self.subset_select(
              self.head_samples,
              topk=self.num_heads,
          )
      elif self.select_strategy == "group":
          self.subset_heads, self.subset_weights = self.group_selet(
              self.head_samples,
              topk=self.num_heads,
          )
      else:
          raise ValueError(
              "{} is not supported".format(self.select_strategy)
          )
  
      self.batch_subset = self.subset_heads[task_ids, :, :]
      self.batch_weights = self.subset_weights[task_ids, :, :]
  
file_path: fairseq/examples/attention_head_selection/src/modules/attn_head_selector.py
imports: import torch
  import torch.nn
  import math
label_desc: Use image node_id 5 for calling the AttnHeadSelector obj's underlying member method code with example usage: obj.head_select(task_ids) without return types