code: def setup_initial_state(
      constructor: Callable[
          [TransformerConfig, jax.Array], TransformerLM
      ],
      tx,
      config: TransformerConfig,
      rng: jax.Array,
      mesh: jax.sharding.Mesh,
  ) -> tuple[TrainState, TrainState]:
      """We initialize the model and optimizer state, and optionally load from a
      checkpoint as necessary.
  
      Args:
        constructor: the model constructor
        tx: the optax.GradientTransformation
        config: config object
        rng: jax.prng key
        mesh: jax.devices() mesh
  
      Returns:
        state: the initialized train state
        state_mesh_annotations: the mesh annotations for the train state
      """
  
      # Initialization
  
      with mesh:
          model = constructor(config, rng)
          params, static = model.split(nnx.Param)
          state = TrainState.create(
              apply_fn=static.apply,
              params=params,
              tx=tx,
              graphdef=static,
          )
          state = jax.tree_map(_to_array, state)
          state_spec = nnx.get_partition_spec(state)
          state = jax.lax.with_sharding_constraint(state, state_spec)
  
      state_sharding = nnx.get_named_sharding(state, mesh)
      return state, state_sharding
  
file_path: flax/flax/experimental/nnx/examples/lm1b/utils.py
imports: import logging
  from typing import Any, Callable
  from typing_extensions import Protocol, runtime_checkable
  import jax
  import jax.numpy
  import numpy
  from jax.experimental import mesh_utils
  from flax.experimental.nnx.examples.lm1b.configs import default
  from models import TransformerConfig, TransformerLM
  from flax.experimental import nnx
  from flax.training import train_state
label_desc: Use image node_id 4 for calling a global function with example usage: setup_initial_state(constructor, tx, config, rng, mesh) and returns: state,  state_sharding