code: def _msgpack_ext_pack(x):
      """Messagepack encoders for custom types."""
      # TODO(flax-dev): Array here only work when they are fully addressable.
      # If they are not fully addressable, use the GDA path for checkpointing.
      if isinstance(x, (np.ndarray, jax.Array)):
          return msgpack.ExtType(
              _MsgpackExtType.ndarray, _ndarray_to_bytes(x)
          )
      if np.issctype(type(x)):
          # pack scalar as ndarray
          return msgpack.ExtType(
              _MsgpackExtType.npscalar, _ndarray_to_bytes(np.asarray(x))
          )
      elif isinstance(x, complex):
          return msgpack.ExtType(
              _MsgpackExtType.native_complex,
              msgpack.packb((x.real, x.imag)),
          )
      return x
  
imports: import enum
  import threading
  from contextlib import contextmanager
  from typing import Any, Dict, List
  import jax
  import msgpack
  import numpy
label_desc: Use image node_id 16 for calling a global function with example usage: _msgpack_ext_pack(x) and returns: x, msgpack, msgpack, msgpack