code: def _state_flatten_with_keys(x: State):
      items = sorted(x._mapping.items(), key=lambda item: item[0])
      children = tuple(
          (jtu.DictKey(key), value) for key, value in items
      )
      return children, tuple(x._mapping.keys())
  
imports: from __future__ import annotations
  import typing
  import jax
  import jax.tree_util
  from flax import traverse_util
  from flax.experimental.nnx.nnx import filterlib, reprlib
  from flax.experimental.nnx.nnx.variables import Variable
label_desc: Use image node_id 1 for calling a global function with example usage: _state_flatten_with_keys(x) and returns: children,  tuple