class_docstr: Implements a query dependent temporal localization head.
  
  The boundary regression takes into account both video and query information.
class_docstr_tok: ["Implements","a","query","dependent","temporal","localization","head",".","The","boundary","regression","takes","into","account","both","video","and","query","information","."]
docstr: Builds a query dependent localization head.
  
  Different from LocalizationHead, we assume video_tokens contain information
  from both the input video and text.
  
  Args:
    video_tokens: A 4D float tensor of shape (batch_size, num_texts,
      num_frames, channels) representing the fused video-text tokens.
    text_emb: Not used. A 3D float tensor of shape (batch_size, num_texts,
      channels) representing the text CLS token for each class. The second
      dimension is batch_size * max_num_captions if task=`moment_retrieval`.
      Otherwise, it is num_classes.
    task: Task name. 'temporal_localization', 'highlight_detection' or
      'moment_retrieval'.
    train: Whether or not it is in training.
  
  Returns:
    logits: In the case of 'temporal_localization' or 'highlight_detection',
      `logits` is a 3D float tensor in shape (batch_size, num_frames,
      num_classes * 3) if self.output_per_class_displacements = True,
      otherwise in shape (batch_size, num_frames, num_classes + 2). In the
      case of 'moment_retrieval', `logits` is a 4D tensor in shape
      (batch_size, num_texts, num_frames, 3) storing the logits of frame
      labels and the predicted distances to the start/end times.
docstr_tok: ["Builds","a","query","dependent","localization","head",".","Different","from","LocalizationHead",",","we","assume","video_tokens","contain","information","from","both","the","input","video","and","text",".","Args",":","video_tokens",":","A","4D","float","tensor","of","shape","(","batch_size",",","num_texts",",","num_frames",",","channels",")","representing","the","fused","video-text","tokens",".","text_emb",":","Not","used",".","A","3D","float","tensor","of","shape","(","batch_size",",","num_texts",",","channels",")","representing","the","text","CLS","token","for","each","class",".","The","second","dimension","is","batch_size","*","max_num_captions","if","task=","`","moment_retrieval","`",".","Otherwise",",","it","is","num_classes",".","task",":","Task","name",".","'temporal_localization","'",",","'highlight_detection","'","or","'moment_retrieval","'",".","train",":","Whether","or","not","it","is","in","training",".","Returns",":","logits",":","In","the","case","of","'temporal_localization","'","or","'highlight_detection","'",",","`","logits","`","is","a","3D","float","tensor","in","shape","(","batch_size",",","num_frames",",","num_classes","*","3",")","if","self.output_per_class_displacements","=","True",",","otherwise","in","shape","(","batch_size",",","num_frames",",","num_classes","+","2",")",".","In","the","case","of","'moment_retrieval","'",",","`","logits","`","is","a","4D","tensor","in","shape","(","batch_size",",","num_texts",",","num_frames",",","3",")","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 a query dependent localization head.
  
      Different from LocalizationHead, we assume video_tokens contain information
      from both the input video and text.
  
      Args:
        video_tokens: A 4D float tensor of shape (batch_size, num_texts,
          num_frames, channels) representing the fused video-text tokens.
        text_emb: Not used. A 3D float tensor of shape (batch_size, num_texts,
          channels) representing the text CLS token for each class. The second
          dimension is batch_size * max_num_captions if task=`moment_retrieval`.
          Otherwise, it is num_classes.
        task: Task name. 'temporal_localization', 'highlight_detection' or
          'moment_retrieval'.
        train: Whether or not it is in training.
  
      Returns:
        logits: In the case of 'temporal_localization' or 'highlight_detection',
          `logits` is a 3D float tensor in shape (batch_size, num_frames,
          num_classes * 3) if self.output_per_class_displacements = True,
          otherwise in shape (batch_size, num_frames, num_classes + 2). In the
          case of 'moment_retrieval', `logits` is a 4D tensor in shape
          (batch_size, num_texts, num_frames, 3) storing the logits of frame
          labels and the predicted distances to the start/end times.
      """
      assert video_tokens.ndim == 4 and text_emb.ndim == 3
      bs, _, num_frames, _ = video_tokens.shape
  
      if self.weight_sharing:
          (
              all_classification_logits,
              all_distances,
          ) = self._build_with_weight_sharing(
              video_tokens,
              classification_output_dim=1,
              regression_output_dim=2,
              pyramid_feature_axis=2,
          )
      else:
          (
              all_classification_logits,
              all_distances,
          ) = self._build_without_weight_sharing(
              video_tokens,
              classification_output_dim=1,
              regression_output_dim=2,
              pyramid_feature_axis=2,
          )
  
      all_classification_logits = jnp.concatenate(
          all_classification_logits, axis=2
      )
      if self.output_per_class_displacements:
          all_distances = jnp.concatenate(all_distances, axis=2)
      else:
          all_distances = jnp.concatenate(all_distances, axis=1)
      if task == "moment_retrieval":
          return jnp.concatenate(
              [all_classification_logits, all_distances], axis=-1
          )
      elif (
          task == "temporal_localization"
          or task == "highlight_detection"
      ):
          if self.output_per_class_displacements:
              logits = jnp.concatenate(
                  [all_classification_logits, all_distances], axis=-1
              )
              # Transpose logits into (batch_size, num_frames, num_classes, 3)
              logits = jnp.transpose(logits, (0, 2, 1, 3))
              return logits.reshape((bs, num_frames, -1))
          else:
              # Changes logits shape into (batch_size, num_frames, num_classes).
              all_classification_logits = jnp.transpose(
                  jnp.squeeze(all_classification_logits, axis=-1),
                  [0, 2, 1],
              )
              return jnp.concatenate(
                  [all_classification_logits, all_distances], axis=-1
              )
      else:
          raise ValueError(f"Unsupported task: {task}.")
  
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 1 for calling the QueryDependentLocalizationHead obj's underlying member method code with example usage: obj.__call__(video_tokens, text_emb, task, train) and returns: jnp, logits, jnp