code: def _pred_probs(predictions):
      """Adds the predicted probabilities to prediction dict."""
      # Compute class probabilities on device, needed for Hungarian matching.
      # Does not remove class logits as they are needed for stable loss computation
      # with cross entropy losses.
      predictions["pred_probs"] = nn.softmax(
          predictions["pred_logits"], axis=-1
      )
      for aux_preds in predictions.get("aux_outputs", []):
          aux_preds["pred_probs"] = nn.softmax(
              aux_preds["pred_logits"], axis=-1
          )
  
imports: import collections
  from concurrent import futures
  import functools
  import random
  import time
  from typing import Any, Dict, Tuple, Optional
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  import flax.linen
  from flax.training.checkpoints import restore_checkpoint
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.baselines.detr import train_utils
  from scenic.projects.layout_denoise import train_utils
  from scenic.projects.layout_denoise.datasets import dataset
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 1 for calling a global function with example usage: _pred_probs(predictions) without return types