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: Normalizes predicted distances to the start and end times.
docstr_tok: ["Normalizes","predicted","distances","to","the","start","and","end","times","."]
code: def _normalize_distance(self, distances: jnp.ndarray) -> jnp.ndarray:
      """Normalizes predicted distances to the start and end times."""
      if self.distance_normalizer == "relu_clip":
          distances = nn.relu(distances)
          # We normalize the distances to be 0 and 1.
          return jnp.clip(distances, a_max=1.0)
      elif self.distance_normalizer == "sigmoid":
          return nn.sigmoid(distances)
      elif self.distance_normalizer == "relu":
          return nn.relu(distances)
      else:
          raise ValueError(
              f"Unknown distance_normalizer: {self.distance_normalizer}."
          )
  
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 LocalizationHead obj's underlying member method code with example usage: obj._normalize_distance(distances) and returns: jnp, nn, nn