code: def create_train_state(
      rng,
      config: ml_collections.ConfigDict,
      model,
      image_size,
      learning_rate_fn,
  ):
      """Create initial training state."""
      dynamic_scale = None
      platform = jax.local_devices()[0].platform
      if config.half_precision and platform == "gpu":
          dynamic_scale = dynamic_scale_lib.DynamicScale()
      else:
          dynamic_scale = None
  
      params, batch_stats = initialized(rng, image_size, model)
      tx = optax.sgd(
          learning_rate=learning_rate_fn,
          momentum=config.momentum,
          nesterov=True,
      )
      state = TrainState.create(
          apply_fn=model.apply,
          params=params,
          tx=tx,
          batch_stats=batch_stats,
          dynamic_scale=dynamic_scale,
      )
      return state
  
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 13 for calling a global function with example usage: create_train_state(rng, config, model, image_size, learning_rate_fn) and returns: state