code: def save_model(
      name: Tag | str,
      module: nn.Module,
      state: dict[str, t.Any]
      | FrozenDict[str, t.Any]
      | struct.PyTreeNode,
      *,
      signatures: ModelSignaturesType | None = None,
      labels: dict[str, str] | None = None,
      custom_objects: dict[str, t.Any] | None = None,
      external_modules: t.List[ModuleType] | None = None,
      metadata: dict[str, t.Any] | None = None,
  ) -> bentoml.Model:
      """
      Save a ``flax.linen.Module`` model instance to the BentoML model store.
  
      Args:
          name: The name to give to the model in the BentoML store. This must be a valid
                :obj:`~bentoml.Tag` name.
          module: ``flax.linen.Module`` to be saved.
          signatures: Signatures of predict methods to be used. If not provided, the signatures default to
                      ``predict``. See :obj:`~bentoml.types.ModelSignature` for more details.
          labels: A default set of management labels to be associated with the model. An example is ``{"training-set": "data-1"}``.
          custom_objects: Custom objects to be saved with the model. An example is ``{"my-normalizer": normalizer}``.
                          Custom objects are currently serialized with cloudpickle, but this implementation is subject to change.
          external_modules: user-defined additional python modules to be saved alongside the model or custom objects,
                            e.g. a tokenizer module, preprocessor module, model configuration module.
          metadata: Metadata to be associated with the model. An example is ``{"bias": 4}``.
                    Metadata is intended for display in a model management UI and therefore must be a
                    default Python type, such as :obj:`str` or :obj:`int`.
  
      Returns:
          :obj:`~bentoml.Tag`: A tag that can be used to access the saved model from the BentoML model store.
  
      Example:
  
      .. code-block:: python
  
         import jax
  
         rng, init_rng = jax.random.split(rng)
         state = create_train_state(init_rng, config)
  
         for epoch in range(1, config.num_epochs + 1):
             rng, input_rng = jax.random.split(rng)
             state, train_loss, train_accuracy = train_epoch(
                 state, train_ds, config.batch_size, input_rng
             )
             _, test_loss, test_accuracy = apply_model(
                 state, test_ds["image"], test_ds["label"]
             )
  
             logger.info(
                 "epoch:% 3d, train_loss: %.4f, train_accuracy: %.2f, test_loss: %.4f, test_accuracy: %.2f",
                 epoch, train_loss, train_accuracy * 100, test_loss, test_accuracy * 100
             )
  
         # `Save` the model with BentoML
         tag = bentoml.flax.save_model("mnist", CNN(), state)
      """
      if not isinstance(module, nn.Module):
          raise BentoMLException(
              f"'bentoml.flax.save_model()' only support saving 'flax.linen.Module' object. Got {module.__class__.__name__} instead."
          )
  
      context = ModelContext(
          framework_name="flax",
          framework_versions={
              "flax": get_pkg_version("flax"),
              "jax": get_pkg_version("jax"),
              "jaxlib": get_pkg_version("jaxlib"),
          },
      )
  
      if signatures is None:
          signatures = {"__call__": {"batchable": False}}
          logger.info(
              'Using the default model signature for Flax (%s) for model "%s".',
              signatures,
              name,
          )
      custom_objects = {} if custom_objects is None else custom_objects
      custom_objects["_module"] = module
  
      with bentoml.models.create(
          name,
          module=MODULE_NAME,
          api_version=API_VERSION,
          signatures=signatures,
          labels=labels,
          options=ModelOptions(),
          custom_objects=custom_objects,
          external_modules=external_modules,
          metadata=metadata,
          context=context,
      ) as bento_model:
          with open(bento_model.path_of(MODEL_FILENAME), "wb") as f:
              f.write(serialization.to_bytes(state))
          return bento_model
  
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 3 for calling a global function with example usage: save_model(name, module, state) and returns: bento_model