code: def load_model(
      bento_model: str | Tag | bentoml.Model,
      init: bool = True,
      device: str | XlaBackend = "cpu",
  ) -> tuple[nn.Module, dict[str, t.Any]]:
      """
      Load the ``flax.linen.Module`` model instance with the given tag from the local BentoML model store.
  
      Args:
          bento_model: Either the tag of the model to get from the store, or a BentoML `~bentoml.Model` instance to load the model from.
          init: Whether to initialize the state dict of given ``flax.linen.Module``.
                By default, the weights and values will be put to ``jnp.ndarray``. If ``init`` is set to ``False``,
                The state_dict will only be put to given accelerator device instead.
          device: The device to put the state dict to. By default, it will be put to ``cpu``. This is
                  only used when ``init`` is set to ``False``.
  
      Returns:
          A tuple of ``flax.linen.Module`` as well as its ``state_dict`` from the model store.
  
      Example:
  
      .. code-block:: python
  
         import bentoml
         import jax
  
         net, state_dict = bentoml.flax.load_model("mnist:latest")
         predict_fn = jax.jit(lambda s: net.apply({"params": state_dict["params"]}, x))
         results = predict_fn(jnp.ones((1, 28, 28, 1)))
      """
      # NOTE: we need to hide all GPU from TensorFlow, otherwise it will try to allocate
      # memory on the GPU and make it unavailable for JAX.
      tf.config.experimental.set_visible_devices([], "GPU")
  
      if not isinstance(bento_model, bentoml.Model):
          bento_model = get(bento_model)
      if bento_model.info.module not in (MODULE_NAME, __name__):
          raise NotFound(
              f"Model {bento_model.tag} was saved with module {bento_model.info.module}, failed loading with {MODULE_NAME}."
          )
      if "_module" not in bento_model.custom_objects:
          raise BentoMLException(
              f"Model {bento_model.tag} was either corrupt or not saved with 'bentoml.flax.save_model()'."
          )
      module: nn.Module = bento_model.custom_objects["_module"]
  
      serialized = bento_model.path_of(MODEL_FILENAME)
      try:
          with open(serialized, "rb") as f:
              state_dict: dict[str, t.Any] = serialization.from_bytes(
                  module, f.read()
              )
      except (
          UnpicklingError,
          msgpack.exceptions.ExtraData,
          UnicodeDecodeError,
      ) as err:
          raise BentoMLException(
              f"Unable to covert model {bento_model.tag}'s state_dict: {err}"
          ) from None
  
      # ensure that all arrays are restored as jnp.ndarray
      # NOTE: This is to prevent a bug this will be fixed in Flax >= v0.3.4:
      # https://github.com/google/flax/issues/1261
      if init:
          state_dict = jax.tree_util.tree_map(jnp.array, state_dict)
      else:
          # keep the params on given device if we don't want to initialize
          state_dict = jax.tree_util.tree_map(
              lambda s: jax.device_put(s, jax.devices(device)[0]),
              state_dict,
          )
      return module, state_dict
  
imports: from __future__ import annotations
  import functools
  import logging
  import typing
  from pickle import UnpicklingError
  from types import ModuleType
  from typing import TYPE_CHECKING
  import msgpack.exceptions
  import bentoml
  from ...exceptions import BentoMLException
  from ...exceptions import MissingDependencyException
  from ...exceptions import NotFound
  from ..models.model import ModelContext
  from ..models.model import PartialKwargsModelOptions
  from ..runner.utils import Params
  from ..types import LazyType
  from ..utils import LazyLoader
  from ..utils.pkg import get_pkg_version
  from .common.jax import JaxArrayContainer
  from .common.jax import jax
  from .common.jax import jnp
label_desc: Use image node_id 2 for calling a global function with example usage: load_model(bento_model, init, device) and returns: module,  state_dict