code: def unbox_logicallypartioned_trainstate(
      boxed_train_state: train_state.TrainState,
  ):
      """Unboxes the flax.LogicallyPartitioned pieces in a train state.
  
      Args:
        boxed_train_state: a train state that includes LogicallyPartitioned
          leaves.
      Returns:
        a TrainState where all all LogicallyPartitioned leaves have been unboxed.
      """
      return jax.tree_util.tree_map(
          lambda x: x.unbox()
          if isinstance(x, nn.spmd.LogicallyPartitioned)
          else x,
          boxed_train_state,
          is_leaf=lambda k: isinstance(k, nn.spmd.LogicallyPartitioned),
      )
  
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 3 for calling a global function with example usage: unbox_logicallypartioned_trainstate(boxed_train_state) and returns: jax