code: def create_train_state(rng):
      quick_draw = QuickDraw()
      params = quick_draw.init(rng, jnp.ones((1, 28, 28, 1)))["params"]
      tx = optax.adam(learning_rate=0.001, b1=0.9, b2=0.999)
      return train_state.TrainState.create(
          apply_fn=quick_draw.apply, params=params, tx=tx
      )
  
imports: from absl import app
  from absl import flags
  import os
  import time
  from flax import linen
  from flax.training import train_state
  import jax
  from jax import lax
  from jax import numpy
  import optax
  import numpy
  import tensorflow
  import tensorflowjs
  import input_pipeline
label_desc: Use image node_id 4 for calling a global function with example usage: create_train_state(rng) and returns: train_state