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: Make TrainState a subscriptable object.
docstr_tok: ["Make","TrainState","a","subscriptable","object","."]
code: def __getitem__(self, item):
      """Make TrainState a subscriptable object."""
      return getattr(self, item)
  
imports: import os
  from typing import Any, Dict, Tuple, Optional
  import flax
  from flax import jax_utils
  from flax import struct
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import optax
  from scenic.dataset_lib import dataset_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 1 for calling the TrainState obj's underlying member method code with example usage: obj.__getitem__(item) and returns: getattr