code: def network(arch: str):
      if arch == "cnn32-3-max":
          return functools.partial(
              convnet.ConvNet,
              scales=3,
              filters=32,
              filters_max=1024,
              pooling=objax.functional.max_pool_2d,
          )
      elif arch == "cnn32-3-mean":
          return functools.partial(
              convnet.ConvNet,
              scales=3,
              filters=32,
              filters_max=1024,
              pooling=objax.functional.average_pool_2d,
          )
      elif arch == "cnn64-3-max":
          return functools.partial(
              convnet.ConvNet,
              scales=3,
              filters=64,
              filters_max=1024,
              pooling=objax.functional.max_pool_2d,
          )
      elif arch == "cnn64-3-mean":
          return functools.partial(
              convnet.ConvNet,
              scales=3,
              filters=64,
              filters_max=1024,
              pooling=objax.functional.average_pool_2d,
          )
      elif arch == "wrn28-1":
          return functools.partial(
              wide_resnet.WideResNet, depth=28, width=1
          )
      elif arch == "wrn28-2":
          return functools.partial(
              wide_resnet.WideResNet, depth=28, width=2
          )
      elif arch == "wrn28-10":
          return functools.partial(
              wide_resnet.WideResNet, depth=28, width=10
          )
      raise ValueError("Architecture not recognized", arch)
  
imports: import functools
  import os
  import shutil
  from typing import Callable
  import json
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
  from absl import app, flags
  import objax
  from objax.jaxboard import SummaryWriter, Summary
  from objax.util import EasyDict
  from objax.zoo import convnet, wide_resnet
  from dataset import DataSet
label_desc: Use image node_id 2 for calling a global function with example usage: network(arch) and returns: functools, functools, functools, functools, functools, functools, functools