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 TAL head with separate decoders.
docstr_tok: ["Builds","TAL","head","with","separate","decoders","."]
code: def _build_without_weight_sharing(
      self,
      tokens: jnp.ndarray,
      classification_output_dim: int,
      regression_output_dim: int,
      pyramid_feature_axis: int,
  ) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
      """Builds TAL head with separate decoders."""
      if self.feature_pyramid_config is None:
          tokens_per_level = [tokens]
      else:
          # pytype: disable=attribute-error
          tokens_per_level = model_utils.split_pyramid_features(
              tokens,
              self.feature_pyramid_config.num_features_level0,
              len(self.feature_pyramid_config.feature_pyramid_levels),
              self.feature_pyramid_config.feature_pyramid_downsample_stride,
              axis=pyramid_feature_axis,
          )
          # pytype: enable=attribute-error
      all_classification_logits = []
      all_distances = []
      for level, x in enumerate(tokens_per_level):
          classification_logits = ConvBlock(
              num_conv_layers=self.num_conv_layers,
              kernel_size=self.kernel_size,
              output_dim=classification_output_dim,
              init_proj_bias=self.init_classification_head_bias,
              name=f"ClassificationHead_{level}",
          )(x)
          if (
              tokens.ndim == 4
              and not self.output_per_class_displacements
          ):
              x = jnp.mean(x, axis=1)
          distances = ConvBlock(
              num_conv_layers=self.num_conv_layers,
              kernel_size=self.kernel_size,
              output_dim=regression_output_dim,
              init_proj_bias=self.init_regression_head_bias,
              name=f"RegressionHead_{level}",
          )(x)
          distances = self._normalize_distance(distances)
          all_classification_logits.append(classification_logits)
          all_distances.append(distances)
      return all_classification_logits, all_distances
  
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 3 for calling the LocalizationHead obj's underlying member method code with example usage: obj._build_without_weight_sharing(tokens, classification_output_dim, regression_output_dim, pyramid_feature_axis) and returns: all_classification_logits,  all_distances