code: def init_train_state(model, tx, config, key):
      """
      We pass in "static" objects like model, tx, config as JAX compares them by
      object hash, and instantiating them inside causes pjit top-level annotations
      to fail to match as pytree prefixes if we re-instantiate.
  
      Args: model, tx, config, key
      """
      input_shape = (
          config.per_device_batch_size,
          config.max_target_length,
      )
      initial_variables = jax.jit(model.init)(
          key, jnp.ones(input_shape, jnp.float32)
      )
  
      state = train_state.TrainState.create(
          apply_fn=model.apply,
          params=initial_variables["params"],
          tx=tx,
      )
      return state
  
file_path: flax/examples/lm1b/utils.py
imports: import functools
  import logging
  import numpy
  import flax.linen
  from flax.linen import partitioning
  from flax.training import train_state
  import jax
  import jax.numpy
  from jax.experimental import mesh_utils
label_desc: Use image node_id 4 for calling a global function with example usage: init_train_state(model, tx, config, key) and returns: state