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
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)))
src/bentoml/_internal/frameworks/flax.py