code: def get_segments_from_frame_predictions_jax(
      class_probs: jnp.ndarray,
      displacements: jnp.ndarray,
      input_mask: jnp.ndarray,
      total_frames: int,
      stride: int,
      sampling_strategy: str = "random",
      displacement_normalizer: str = "duration",
      secs_per_timestep: float = 1.0,
      score_threshold: float = float("-inf"),
      feature_pyramid_config: Optional[
          ml_collections.ConfigDict
      ] = None,
  ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
      """Computes predicted segments based on frame predictions in jnp format.
  
      We assume that all frame predictions are from one video.
  
      Args:
        class_probs: A float array representing the class probabilities in shape
          (num_frames, num_classes).
        displacements: A float array representing the start/end time displacements
          in shape (num_frames, num_classes, 2). All values are greater or equal to
          zero.
        input_mask: A int array representing the input mask in shape (num_frames,).
        total_frames: The total number of frames in this video.
        stride: Temporal stride used in sampling the frames
        sampling_strategy: 'random' or 'linspace'.
        displacement_normalizer: 'none', 'duration', or 'sampled_span'.
        secs_per_timestep: Separation in seconds between two consecutive timesteps.
        score_threshold: Score threshold to filter the detections.
        feature_pyramid_config: Feature pyramid configs.
  
      Returns:
        Class indices in shape (num_frames,), padded ones are filled with -1.
        Class probabilities in shape (num_frames,), padded ones are filled with -1.
        Predicted segment start and end times in shape (num_frames, 2), padded ones
        are filled with -1.
      """
      assert sampling_strategy in {"random", "linspace"}
  
      if feature_pyramid_config is None:
          num_frames = class_probs.shape[0]
      else:
          num_frames = feature_pyramid_config.num_features_level0
      num_classes = class_probs.shape[-1]
      if sampling_strategy == "random":
          # The default setting is to sample the center clip at test time.
          start_frame_offset = jnp.maximum(
              0, (total_frames - num_frames * stride) // 2
          ).astype(jnp.float32)
      else:  # 'linspace'
          start_frame_offset = 0.0
      if displacement_normalizer == "duration":
          displacement_multiplier = total_frames
      elif displacement_normalizer == "sampled_span":
          displacement_multiplier = num_frames * stride
      else:
          displacement_multiplier = 1
      displacements = displacements.at[..., 0].multiply(-1.0)
  
      if feature_pyramid_config is not None:
          feature_pyramid_downsample_stride = (
              feature_pyramid_config.feature_pyramid_downsample_stride
          )
          displacements_per_level = model_utils.split_pyramid_features(
              displacements,
              feature_pyramid_config.num_features_level0,
              len(feature_pyramid_config.feature_pyramid_levels),
              feature_pyramid_downsample_stride,
              axis=0,
              np_backend=jnp,
          )
          normalize_displacements_by_downsample_stride = (
              feature_pyramid_config.get(
                  "normalize_displacements_by_downsample_stride", False
              )
          )
      else:
          displacements_per_level = [displacements]
          feature_pyramid_downsample_stride = 1
          normalize_displacements_by_downsample_stride = False
  
      segments = []
      linspace_frame_indices = jnp.arange(
          0, num_frames, dtype=jnp.float32
      )
      for level, cur_displacements in enumerate(
          displacements_per_level
      ):
          cur_downsample_stride = (
              feature_pyramid_downsample_stride**level
          )
          if normalize_displacements_by_downsample_stride:
              cur_displacements *= cur_downsample_stride
          cur_stride = stride * cur_downsample_stride
          if sampling_strategy == "random":
              frame_indices = jnp.arange(
                  0, num_frames * stride, cur_stride, dtype=jnp.float32
              )[:, None, None]
              cur_segments = (
                  frame_indices
                  + cur_displacements * displacement_multiplier
                  + start_frame_offset
              )
          else:  # 'linspace'
              frame_indices = linspace_frame_indices[
                  jnp.arange(0, num_frames, cur_downsample_stride)
              ][:, None, None]
              cur_segments = (
                  (frame_indices + cur_displacements)
                  * (total_frames - 1)
                  / (num_frames - 1)
              )
          segments.append(cur_segments)
      segments = jnp.concatenate(segments, axis=0)
      input_mask = jnp.array(input_mask, dtype=bool)
      total_frames = jnp.array(
          jnp.full((segments.shape[0], num_classes), total_frames)
      )
      segments = segments.at[..., 0].set(
          jnp.maximum(segments[..., 0], 0)
      )
      segments = segments.at[..., 1].set(
          jnp.minimum(segments[..., 1], total_frames)
      )
      segments = segments * secs_per_timestep
      class_probs = jnp.where(input_mask[:, None], class_probs, -1)
      mask = class_probs >= score_threshold
      inds, class_indices = mask.nonzero(
          size=mask.shape[0], fill_value=-1
      )
      class_probs = jnp.where(
          inds >= 0,
          jnp.take_along_axis(class_probs[:, 0], inds, axis=0),
          -1,
      )
      segments = jnp.where(
          inds[:, None] >= 0,
          jnp.take_along_axis(segments[:, 0], inds[:, None], axis=0),
          -1,
      )
  
      return class_indices, class_probs, segments
  
file_path: scenic/scenic/projects/unloc/postprocessing_utils.py
imports: from typing import Any, List, Optional, Tuple, Union
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.unloc import model_utils
  import tensorflow
label_desc: Use image node_id 8 for calling a global function with example usage: get_segments_from_frame_predictions_jax(class_probs, displacements, input_mask, total_frames, stride, sampling_strategy, displacement_normalizer, secs_per_timestep, score_threshold, feature_pyramid_config) and returns: class_indices,  class_probs,  segments