class_docstr: Single-head Attention layer for label-specific representations
class_docstr_tok: ["Single-head","Attention","layer","for","label-specific","representations"]
code: def pad_and_rearrange(self, q_s, k_s, v_s, batch_idxs):
      # Input is padded representation: n_head x len_inp x d
      # Output is packed representation: (n_head * mb_size) x len_padded x d
      # (along with masks for the attention and output)
      n_head = self.d_l
      d_k, d_v = self.d_k, self.d_v
  
      len_padded = batch_idxs.max_len
      mb_size = batch_idxs.batch_size
      if self.q_as_matrix:
          q_padded = q_s.new_zeros((n_head, mb_size, len_padded, d_k))
      else:
          q_padded = q_s.repeat(
              mb_size, 1, 1
          )  # (d_l * mb_size) x 1 x d_k
      k_padded = k_s.new_zeros((n_head, mb_size, len_padded, d_k))
      v_padded = v_s.new_zeros((n_head, mb_size, len_padded, d_v))
      invalid_mask = q_s.new_ones((mb_size, len_padded), dtype=DTYPE)
  
      for i, (start, end) in enumerate(
          zip(
              batch_idxs.boundaries_np[:-1],
              batch_idxs.boundaries_np[1:],
          )
      ):
          if self.q_as_matrix:
              q_padded[:, i, : end - start, :] = q_s[:, start:end, :]
          k_padded[:, i, : end - start, :] = k_s[:, start:end, :]
          v_padded[:, i, : end - start, :] = v_s[:, start:end, :]
          invalid_mask[i, : end - start].fill_(False)
  
      if self.q_as_matrix:
          q_padded = q_padded.view(-1, len_padded, d_k)
          attn_mask = (
              invalid_mask.unsqueeze(1)
              .expand(mb_size, len_padded, len_padded)
              .repeat(n_head, 1, 1)
          )
      else:
          attn_mask = invalid_mask.unsqueeze(1).repeat(n_head, 1, 1)
  
      output_mask = (~invalid_mask).repeat(n_head, 1)
  
      return (
          q_padded,
          k_padded.view(-1, len_padded, d_k),
          v_padded.view(-1, len_padded, d_v),
          attn_mask,
          output_mask,
      )
  
imports: import numpy
  import functools
  import sys
  import torch
  from torch.autograd import Variable
  import torch.nn
  import torch.nn.init
label_desc: Use image node_id 3 for calling the LabelAttention obj's underlying member method code with example usage: obj.pad_and_rearrange(q_s, k_s, v_s, batch_idxs) and returns: q_padded,  k_padded,  v_padded,  attn_mask,  output_mask