code: def re_add_axis_names(
      train_state: TrainState,
      params_axes: Any,
      axes_name: str = "params_axes",
  ) -> TrainState:
      """Adds axes_names to model_state for a train state.
  
      Args:
        train_state: Training state.
        params_axes: Model axes metadata to re-add.
        axes_name: the string specifying the name in the model_state
  
      Returns:
        New train state without axes_names in model_state, axes_names metadata if it
        was removed (so it can be re-added).
      """
      if params_axes:
          model_state = frozen_dict.unfreeze(train_state.model_state)
          model_state[axes_name] = params_axes
          return train_state.replace(
              model_state=frozen_dict.freeze(model_state)
          )
      else:
          return train_state
  
imports: import dataclasses
  import operator
  import re
  from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
  from absl import logging
  from big_vision import utils
  from big_vision.models import vit
  from clu import metric_writers
  import flax
  from flax import jax_utils
  from flax.core import frozen_dict
  import jax
  import jax.numpy
  import ml_collections
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.projects.knowledge_visual_language.models import constants
  from scenic.projects.t5 import model
  from scenic.train_lib import train_utils
  from scenic.train_lib.train_utils import TrainState
label_desc: Use image node_id 7 for calling a global function with example usage: re_add_axis_names(train_state, params_axes, axes_name) and returns: train_state, train_state