code: def re_add_axis_names(
      train_state: OptaxTrainState,
      param_axes: Any,
      axes_name: str = "param_axes",
  ) -> OptaxTrainState:
      """Adds axes_names to model_state for a train state.
  
      Args:
        train_state: Training state.
        param_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 param_axes:
          model_state = frozen_dict.unfreeze(train_state.model_state)
          model_state[axes_name] = param_axes
          return train_state.replace(
              model_state=frozen_dict.freeze(model_state)
          )
      else:
          return train_state
  
imports: import os
  import re
  from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
  from absl import logging
  import flax
  from flax import struct
  from flax.core import frozen_dict
  import flax.linen
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.baselines.clip import model
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 11 for calling a global function with example usage: re_add_axis_names(train_state, param_axes, axes_name) and returns: train_state, train_state