code: def pop_axes_names(
      train_state: OptaxTrainState, axes_name: str = "param_axes"
  ) -> Tuple[OptaxTrainState, Optional[Any]]:
      """Removes axes_names from model_state for a train state.
  
      Args:
        train_state: Training state.
        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).
      """
      model_state = train_state.model_state
      if axes_name in train_state.model_state:
          model_state, param_axes = frozen_dict.freeze(model_state).pop(
              axes_name
          )
          return (
              train_state.replace(model_state=model_state),
              param_axes,
          )
      else:
          return train_state, None
  
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 10 for calling a global function with example usage: pop_axes_names(train_state, axes_name) and returns: train_state,  param_axes, train_state,  None