code: def create_train_state(rng, config: ml_collections.ConfigDict, model):
      """Create initial training state."""
      params = get_initial_params(rng, model)
      tx = optax.chain(
          optax.sgd(
              learning_rate=config.learning_rate,
              momentum=config.momentum,
          ),
          optax.add_decayed_weights(weight_decay=config.weight_decay),
      )
      state = TrainState.create(
          apply_fn=model.apply, params=params, tx=tx
      )
      return state
  
imports: from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple, Union
  from absl import logging
  from flax import struct
  from flax.metrics import tensorboard
  from flax.training import train_state
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  import tensorflow
  import input_pipeline
  import models
label_desc: Use image node_id 3 for calling a global function with example usage: create_train_state(rng, config, model) and returns: state