class_docstr: Implements a linear head.
class_docstr_tok: ["Implements","a","linear","head","."]
docstr: Builds a linear head.
  
  Args:
    video_tokens: A ND float tensor of shape (batch_size, ..., channels)
      representing the video tokens.
    text_tokens: A ND float tensor of shape (batch_size, ..., channels)
      representing the video tokens. Not used.
    task: 'action_segmentation'.
    train: Whether or not the model is under training. Not used.
  
  Returns:
    logits: A (N-1)D float tensor of shape (batch_size, ...) if output_dim ==
    1. Otherwise, A ND float tensor of shape (batch_size, ..., output_dim).
docstr_tok: ["Builds","a","linear","head",".","Args",":","video_tokens",":","A","ND","float","tensor","of","shape","(","batch_size",",","...",",","channels",")","representing","the","video","tokens",".","text_tokens",":","A","ND","float","tensor","of","shape","(","batch_size",",","...",",","channels",")","representing","the","video","tokens",".","Not","used",".","task",":","'action_segmentation","'",".","train",":","Whether","or","not","the","model","is","under","training",".","Not","used",".","Returns",":","logits",":","A","(","N-1",")","D","float","tensor","of","shape","(","batch_size",",","...",")","if","output_dim","==","1",".","Otherwise",",","A","ND","float","tensor","of","shape","(","batch_size",",","...",",","output_dim",")","."]
code: def __call__(
      self,
      video_tokens: jnp.ndarray,
      text_tokens: Optional[jnp.ndarray],
      task: str,
      train: bool,
  ) -> jnp.ndarray:
      """Builds a linear head.
  
      Args:
        video_tokens: A ND float tensor of shape (batch_size, ..., channels)
          representing the video tokens.
        text_tokens: A ND float tensor of shape (batch_size, ..., channels)
          representing the video tokens. Not used.
        task: 'action_segmentation'.
        train: Whether or not the model is under training. Not used.
  
      Returns:
        logits: A (N-1)D float tensor of shape (batch_size, ...) if output_dim ==
        1. Otherwise, A ND float tensor of shape (batch_size, ..., output_dim).
      """
      output = nn.Dense(
          self.output_dim,
          kernel_init=nn.initializers.xavier_uniform(),
          bias_init=jax.nn.initializers.constant(self.init_head_bias),
          name="output_projection",
      )(video_tokens)
      if self.output_dim == 1:
          return jnp.squeeze(output, axis=-1)
      return output
  
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 LinearHead obj's underlying member method code with example usage: obj.__call__(video_tokens, text_tokens, task, train) and returns: output, jnp