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(  # type: ignore
        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
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)
src/bentoml/_internal/frameworks/flax.py