class_docstr: Dataclass to keep track of state of training.
  
  The state of training is structured as a flax.struct.dataclass, which enables
  instances of this class to be passed into jax transformations like tree_map
  and pmap.
class_docstr_tok: ["Dataclass","to","keep","track","of","state","of","training",".","The","state","of","training","is","structured","as","a","flax.struct.dataclass",",","which","enables","instances","of","this","class","to","be","passed","into","jax","transformations","like","tree_map","and","pmap","."]
docstr: Return the value for key if it exists otherwise the default.
code: def get(self, keyname: str, default: Optional[Any] = None) -> Any:
      """Return the value for key if it exists otherwise the default."""
      try:
          return self[keyname]
      except KeyError:
          return default
  
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 2 for calling the OptaxTrainState obj's underlying member method code with example usage: obj.get(keyname, default) and returns: self, default