code: def create_train_state(rng, config):
      """Creates initial `TrainState`."""
      cnn = CNN()
      params = cnn.init(rng, jnp.ones([1, 28, 28, 1]))["params"]
      tx = optax.sgd(config.learning_rate, config.momentum)
      return train_state.TrainState.create(
          apply_fn=cnn.apply, params=params, tx=tx
      )
  
file_path: flax/examples/mnist/train.py
imports: from absl import logging
  from flax import linen
  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_datasets
label_desc: Use image node_id 5 for calling a global function with example usage: create_train_state(rng, config) and returns: train_state