code: def _handle_legacy_format(train_state):
      """Handle legacy format.
  
      To remove this function, make sure all checkpoints that are to be loaded are
      in the non-legacy format:
        * checkpoint['params'] is a PyTree of parameters;
        * checkpoint['model_state']['batch_stats'] is a PyTree of batch statistics.
  
      Args:
        train_state: A train state.
      """
      if "params" not in train_state:
          train_state["params"] = train_state.pop("optimizer").pop(
              "target"
          )
          if "param" in train_state["params"]:
              train_state["params"] = train_state["params"]["param"]
      if (
          "model_state" in train_state
          and "batch_stats" not in train_state["model_state"]
      ):
          bad_restored_batch_stats = train_state.pop("model_state")
          good_restored_batch_stats = {}
          for key, value in bad_restored_batch_stats.items():
              current = good_restored_batch_stats
              subkeys = [subkey for subkey in key.split("/") if subkey]
              for subkey in subkeys[:-1]:
                  if subkey not in current:
                      new_current = {}
                      current[subkey] = new_current
                      current = new_current
                  else:
                      current = current[subkey]
              current[subkeys[-1]] = value
          train_state["model_state"] = {}
          train_state["model_state"][
              "batch_stats"
          ] = good_restored_batch_stats
  
imports: from collections.abc import Mapping
  from concurrent import futures
  import functools
  import time
  from typing import Any, Callable, Dict, Sequence, Tuple, Type, Union
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  import flax
  from flax import jax_utils
  from flax.training.checkpoints import restore_checkpoint
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.baselines.deformable_detr import coco_eval
  from scenic.projects.baselines.deformable_detr import evaluate
  from scenic.projects.baselines.deformable_detr.model import DeformableDETRModel
  from scenic.projects.baselines.detr import train_utils
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
  from scenic.train_lib.train_utils import PyTree
  from scenic.train_lib.train_utils import TrainState
label_desc: Use image node_id 3 for calling a global function with example usage: _handle_legacy_format(train_state) without return types