class_docstr: Implements a localization head.
  
  We follow the head design in ActionFormer: https://arxiv.org/abs/2202.07925.
  In this head, we apply convolutions to generate per-frame logits for labels
  as well as the distance to the start/end times.
  
  Attributes:
    num_conv_layers: Number of convolution layers.
    kernel_size: Convolution kernel size.
    num_classes: Number of classes.
    init_classification_head_bias: Initial bias value for the classification
      head.
    init_regression_head_bias: Initial bias value for the regression head.
    distance_normalizer: We normalize the predicted displacements to be [0, 1].
      Options are: `relu_clip`, `sigmoid`, or 'relu'. When using `relu_clip`,
      the distances are first fed into relu() and then clipped to be [0, 1].
      When using `sigmoid`, the distances are normalized into [0, 1] using the
      `sigmoid` function. When using `relu`, distances are fed into a relu() but
      not normalized.
    weight_sharing: Whether or not to share the weights among decoders from
      different pyramid levels. If True, we will also learn a `scale` and a
      `bias` term for each level.
    feature_pyramid_config: Feature pyramid config.
    output_per_class_displacements: Whether or not to predict start/end time
      displacements for each class.
class_docstr_tok: ["Implements","a","localization","head",".","We","follow","the","head","design","in","ActionFormer",":","https",":","\/\/arxiv.org\/abs\/2202.07925",".","In","this","head",",","we","apply","convolutions","to","generate","per-frame","logits","for","labels","as","well","as","the","distance","to","the","start\/end","times",".","Attributes",":","num_conv_layers",":","Number","of","convolution","layers",".","kernel_size",":","Convolution","kernel","size",".","num_classes",":","Number","of","classes",".","init_classification_head_bias",":","Initial","bias","value","for","the","classification","head",".","init_regression_head_bias",":","Initial","bias","value","for","the","regression","head",".","distance_normalizer",":","We","normalize","the","predicted","displacements","to","be","[","0",",","1","]",".","Options","are",":","`","relu_clip","`",",","`","sigmoid","`",",","or","'relu","'",".","When","using","`","relu_clip","`",",","the","distances","are","first","fed","into","relu","(",")","and","then","clipped","to","be","[","0",",","1","]",".","When","using","`","sigmoid","`",",","the","distances","are","normalized","into","[","0",",","1","]","using","the","`","sigmoid","`","function",".","When","using","`","relu","`",",","distances","are","fed","into","a","relu","(",")","but","not","normalized",".","weight_sharing",":","Whether","or","not","to","share","the","weights","among","decoders","from","different","pyramid","levels",".","If","True",",","we","will","also","learn","a","`","scale","`","and","a","`","bias","`","term","for","each","level",".","feature_pyramid_config",":","Feature","pyramid","config",".","output_per_class_displacements",":","Whether","or","not","to","predict","start\/end","time","displacements","for","each","class","."]
docstr: Builds the temporal localization head.
  
  The head design is following ActionFormer, https://arxiv.org/abs/2202.07925.
  
  Args:
    video_tokens: A 3D float tensor of shape (batch_size, num_frames,
      channels) representing the fused video-text tokens.
    text_emb: Not used.
    task: Task name. Only supports 'temporal_localization' and
      'highlight_detection'.
    train: Whether or not it is in training.
  
  Returns:
    logits: A 3D float tensor in shape (batch_size, num_frames,
      num_classes * 3) if output_per_class_displacements = True, otherwise in
      shape (batch_size, num_frames, num_classes + 2) storing the logits of
      frame labels and the predicted distances to the start/end times.
docstr_tok: ["Builds","the","temporal","localization","head",".","The","head","design","is","following","ActionFormer",",","https",":","\/\/arxiv.org\/abs\/2202.07925",".","Args",":","video_tokens",":","A","3D","float","tensor","of","shape","(","batch_size",",","num_frames",",","channels",")","representing","the","fused","video-text","tokens",".","text_emb",":","Not","used",".","task",":","Task","name",".","Only","supports","'temporal_localization","'","and","'highlight_detection","'",".","train",":","Whether","or","not","it","is","in","training",".","Returns",":","logits",":","A","3D","float","tensor","in","shape","(","batch_size",",","num_frames",",","num_classes","*","3",")","if","output_per_class_displacements","=","True",",","otherwise","in","shape","(","batch_size",",","num_frames",",","num_classes","+","2",")","storing","the","logits","of","frame","labels","and","the","predicted","distances","to","the","start\/end","times","."]
code: def __call__(
      self,
      video_tokens: jnp.ndarray,
      text_emb: jnp.ndarray,
      task: str,
      train: bool,
  ) -> jnp.ndarray:
      """Builds the temporal localization head.
  
      The head design is following ActionFormer, https://arxiv.org/abs/2202.07925.
  
      Args:
        video_tokens: A 3D float tensor of shape (batch_size, num_frames,
          channels) representing the fused video-text tokens.
        text_emb: Not used.
        task: Task name. Only supports 'temporal_localization' and
          'highlight_detection'.
        train: Whether or not it is in training.
  
      Returns:
        logits: A 3D float tensor in shape (batch_size, num_frames,
          num_classes * 3) if output_per_class_displacements = True, otherwise in
          shape (batch_size, num_frames, num_classes + 2) storing the logits of
          frame labels and the predicted distances to the start/end times.
      """
      del text_emb
      assert (
          task == "temporal_localization"
          or task == "highlight_detection"
      )
      assert video_tokens.ndim == 3
      bs, num_frames, _ = video_tokens.shape
      inputs = video_tokens
      if self.output_per_class_displacements:
          regression_output_dim = 2 * self.num_classes
      else:
          regression_output_dim = 2
      if self.weight_sharing:
          (
              all_classification_logits,
              all_distances,
          ) = self._build_with_weight_sharing(
              inputs,
              classification_output_dim=self.num_classes,
              regression_output_dim=regression_output_dim,
              pyramid_feature_axis=1,
          )
      else:
          (
              all_classification_logits,
              all_distances,
          ) = self._build_without_weight_sharing(
              inputs,
              classification_output_dim=self.num_classes,
              regression_output_dim=regression_output_dim,
              pyramid_feature_axis=1,
          )
  
      all_classification_logits = jnp.concatenate(
          all_classification_logits, axis=1
      )
      all_distances = jnp.concatenate(all_distances, axis=1)
      if self.output_per_class_displacements:
          all_classification_logits = all_classification_logits[
              ..., jnp.newaxis
          ]
          all_distances = all_distances.reshape(
              (bs, num_frames, self.num_classes, 2)
          )
      logits = jnp.concatenate(
          [all_classification_logits, all_distances], axis=-1
      )
      return logits.reshape((bs, num_frames, -1))
  
imports: from typing import List, Optional, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.projects.unloc import model_utils
label_desc: Use image node_id 4 for calling the LocalizationHead obj's underlying member method code with example usage: obj.__call__(video_tokens, text_emb, task, train) and returns: logits