code: def forward(self, attn, value, mask):
      heads, output_num_tokens, eps, device, dtype = (
          attn.shape[1],
          self.output_num_tokens,
          self.eps,
          attn.device,
          attn.dtype,
      )
  
      # first get the attention values for CLS token to all other tokens
  
      cls_attn = attn[..., 0, 1:]
  
      # calculate the norms of the values, for weighting the scores, as described in the paper
  
      value_norms = value[..., 1:, :].norm(dim=-1)
  
      # weigh the attention scores by the norm of the values, sum across all heads
  
      cls_attn = einsum("b h n, b h n -> b n", cls_attn, value_norms)
  
      # normalize to 1
  
      normed_cls_attn = cls_attn / (
          cls_attn.sum(dim=-1, keepdim=True) + eps
      )
  
      # instead of using inverse transform sampling, going to invert the softmax and use gumbel-max sampling instead
  
      pseudo_logits = log(normed_cls_attn)
  
      # mask out pseudo logits for gumbel-max sampling
  
      mask_without_cls = mask[:, 1:]
      mask_value = -torch.finfo(attn.dtype).max / 2
      pseudo_logits = pseudo_logits.masked_fill(
          ~mask_without_cls, mask_value
      )
  
      # expand k times, k being the adaptive sampling number
  
      pseudo_logits = repeat(
          pseudo_logits, "b n -> b k n", k=output_num_tokens
      )
      pseudo_logits = pseudo_logits + sample_gumbel(
          pseudo_logits.shape, device=device, dtype=dtype
      )
  
      # gumble-max and add one to reserve 0 for padding / mask
  
      sampled_token_ids = pseudo_logits.argmax(dim=-1) + 1
  
      # calculate unique using torch.unique and then pad the sequence from the right
  
      unique_sampled_token_ids_list = [
          torch.unique(t, sorted=True)
          for t in torch.unbind(sampled_token_ids)
      ]
      unique_sampled_token_ids = pad_sequence(
          unique_sampled_token_ids_list, batch_first=True
      )
  
      # calculate the new mask, based on the padding
  
      new_mask = unique_sampled_token_ids != 0
  
      # CLS token never gets masked out (gets a value of True)
  
      new_mask = F.pad(new_mask, (1, 0), value=True)
  
      # prepend a 0 token id to keep the CLS attention scores
  
      unique_sampled_token_ids = F.pad(
          unique_sampled_token_ids, (1, 0), value=0
      )
      expanded_unique_sampled_token_ids = repeat(
          unique_sampled_token_ids, "b n -> b h n", h=heads
      )
  
      # gather the new attention scores
  
      new_attn = batched_index_select(
          attn, expanded_unique_sampled_token_ids, dim=2
      )
  
      # return the sampled attention scores, new mask (denoting padding), as well as the sampled token indices (for the residual)
      return new_attn, new_mask, unique_sampled_token_ids
  
imports: import torch
  import torch.nn.functional
  from torch.nn.utils.rnn import pad_sequence
  from torch import nn, einsum
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 2 for calling the AdaptiveTokenSampling obj's underlying member method code with example usage: obj.forward(attn, value, mask) and returns: new_attn,  new_mask,  unique_sampled_token_ids