code: def pick_model_class():
      """Picks one of PureJaxMNIST or FlaxMNIST."""
      if _MODEL.value == "mnist_pure_jax":
          return mnist_lib.PureJaxMNIST
      elif _MODEL.value == "mnist_flax":
          return mnist_lib.FlaxMNIST
      else:
          raise ValueError(f"Unrecognized model: {_MODEL.value}")
  
file_path: jax/jax/experimental/jax2tf/examples/saved_model_main.py
imports: import logging
  import os
  from absl import app
  from absl import flags
  from jax.experimental.jax2tf.examples import mnist_lib
  from jax.experimental.jax2tf.examples import saved_model_lib
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: pick_model_class() and returns: mnist_lib, mnist_lib