code: def create_train_state(ctable):
      model = models.Seq2seq(
          teacher_force=False,
          hidden_size=train.FLAGS.hidden_size,
          vocab_size=ctable.vocab_size,
      )
      params = train.get_initial_params(
          model, jax.random.key(0), ctable
      )
      tx = optax.adam(train.FLAGS.learning_rate)
      state = train_state.TrainState.create(
          apply_fn=model.apply, params=params, tx=tx
      )
      return state
  
file_path: flax/examples/seq2seq/train_test.py
imports: import functools
  from absl.testing import absltest
  from flax.training import train_state
  import jax
  from jax import random
  import numpy
  import optax
  import input_pipeline
  import train
  import models
label_desc: Use image node_id 2 for calling a global function with example usage: create_train_state(ctable) and returns: state