code: def create_model(*, model_cls, half_precision, **kwargs):
      platform = jax.local_devices()[0].platform
      if half_precision:
          if platform == "tpu":
              model_dtype = jnp.bfloat16
          else:
              model_dtype = jnp.float16
      else:
          model_dtype = jnp.float32
      return model_cls(
          num_classes=NUM_CLASSES, dtype=model_dtype, **kwargs
      )
  
imports: import functools
  import time
  from typing import Any
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax.training import checkpoints
  from flax.training import common_utils
  from flax.training import dynamic_scale
  from flax.training import train_state
  import jax
  from jax import lax
  import jax.numpy
  from jax import random
  import ml_collections
  import optax
  import tensorflow
  import tensorflow_datasets
  import input_pipeline
  import models
label_desc: Use image node_id 1 for calling a global function with example usage: create_model() and returns: model_cls