code: def save_model(
      model,
      path,
      conda_env=None,
      code_paths=None,
      mlflow_model=None,
      custom_objects=None,
      signature: ModelSignature = None,
      input_example: ModelInputExample = None,
      pip_requirements=None,
      extra_pip_requirements=None,
      saved_model_kwargs=None,
      keras_model_kwargs=None,
      metadata=None,
  ):
      """
      Save a TF2 core model (inheriting tf.Module) or Keras model in MLflow Model format to a path on
      the local file system.
  
      .. note::
          If you save a Keras or TensorFlow model without a signature, inference with
          :py:func:`mlflow.pyfunc.spark_udf()` will not work unless the model's pyfunc
          representation accepts pandas DataFrames as inference inputs.
          You can infer a model's signature by calling the :py:func:`mlflow.models.infer_signature()`
          API on features from the model's test dataset. You can also manually create a model
          signature, for example:
  
          .. code-block:: python
              :caption: Example of creating signature for saving TensorFlow and `tf.Keras` models
  
              from mlflow.types.schema import Schema, TensorSpec
              from mlflow.models import ModelSignature
              import numpy as np
  
              input_schema = Schema(
                  [
                      TensorSpec(np.dtype(np.uint64), (-1, 5), "field1"),
                      TensorSpec(np.dtype(np.float32), (-1, 3, 2), "field2"),
                  ]
              )
              # Create the signature for a model that requires 2 inputs:
              #  - Input with name "field1", shape (-1, 5), type "np.uint64"
              #  - Input with name "field2", shape (-1, 3, 2), type "np.float32"
              signature = ModelSignature(inputs=input_schema)
  
      Args:
          model: The Keras model or Tensorflow module to be saved.
          path: Local path where the MLflow model is to be saved.
          conda_env: {{ conda_env }}
          code_paths: A list of local filesystem paths to Python file dependencies (or directories
              containing file dependencies). These files are *prepended* to the system
              path when the model is loaded.
          mlflow_model: MLflow model configuration to which to add the ``tensorflow`` flavor.
          custom_objects: A Keras ``custom_objects`` dictionary mapping names (strings) to
              custom classes or functions associated with the Keras model. MLflow saves
              these custom layers using CloudPickle and restores them automatically
              when the model is loaded with :py:func:`mlflow.tensorflow.load_model` and
              :py:func:`mlflow.pyfunc.load_model`.
          signature: {{ signature }}
          input_example: {{ input_example }}
          pip_requirements: {{ pip_requirements }}
          extra_pip_requirements: {{ extra_pip_requirements }}
          saved_model_kwargs: a dict of kwargs to pass to ``tensorflow.saved_model.save`` method
              if the model to be saved is a Tensorflow module.
          keras_model_kwargs: a dict of kwargs to pass to ``model.save`` method if the model
              to be saved is a keras model.
          metadata: Custom metadata dictionary passed to the model and stored in the MLmodel file.
  
              .. Note:: Experimental: This parameter may change or be removed in a future
                                      release without warning.
      """
      import tensorflow as tf
      from tensorflow.keras.models import Model as KerasModel
  
      if signature is None and input_example is not None:
          wrapped_model = None
          if isinstance(model, KerasModel):
              wrapped_model = _KerasModelWrapper(model, signature)
          elif isinstance(model, tf.Module):
              wrapped_model = _TF2ModuleWrapper(model, signature)
          if wrapped_model is not None:
              signature = _infer_signature_from_input_example(
                  input_example, wrapped_model
              )
      elif signature is False:
          signature = None
  
      if signature is None:
          _logger.warning(_NO_MODEL_SIGNATURE_WARNING)
      else:
          num_inputs = len(signature.inputs.inputs)
          if num_inputs == 0:
              raise MlflowException(
                  "The model signature's input schema must contain at least one field.",
                  error_code=INVALID_PARAMETER_VALUE,
              )
          for field in signature.inputs.inputs:
              if not isinstance(field, TensorSpec):
                  raise MlflowException(
                      "All fields in the model signature's input schema must be of type TensorSpec.",
                      error_code=INVALID_PARAMETER_VALUE,
                  )
              if field.shape[0] != -1:
                  raise MlflowException(
                      "All fields in the model signature's input schema must have a shape "
                      "in which the first dimension is a variable dimension.",
                      error_code=INVALID_PARAMETER_VALUE,
                  )
  
      _validate_env_arguments(
          conda_env, pip_requirements, extra_pip_requirements
      )
  
      # check if path exists
      path = os.path.abspath(path)
      _validate_and_prepare_target_save_path(path)
  
      code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
  
      if mlflow_model is None:
          mlflow_model = Model()
      if signature is not None:
          mlflow_model.signature = signature
      if input_example is not None:
          _save_example(mlflow_model, input_example, path)
      if metadata is not None:
          mlflow_model.metadata = metadata
  
      if isinstance(model, KerasModel):
          keras_model_kwargs = keras_model_kwargs or {}
  
          data_subpath = "data"
          # construct new data folder in existing path
          data_path = os.path.join(path, data_subpath)
          os.makedirs(data_path)
          model_subpath = os.path.join(data_subpath, _MODEL_SAVE_PATH)
  
          keras_module = importlib.import_module("tensorflow.keras")
          # save custom objects if there are custom objects
          if custom_objects is not None:
              _save_keras_custom_objects(
                  data_path, custom_objects, _CUSTOM_OBJECTS_SAVE_PATH
              )
          # save custom objects stored within _GLOBAL_CUSTOM_OBJECTS
          if global_custom_objects := get_global_custom_objects():
              _save_keras_custom_objects(
                  data_path,
                  global_custom_objects,
                  _GLOBAL_CUSTOM_OBJECTS_SAVE_PATH,
              )
  
          # save keras module spec to path/data/keras_module.txt
          with open(
              os.path.join(data_path, _KERAS_MODULE_SPEC_PATH), "w"
          ) as f:
              f.write(keras_module.__name__)
  
          # Use the SavedModel format if `save_format` is unspecified
          save_format = keras_model_kwargs.get("save_format", "tf")
  
          # save keras save_format to path/data/save_format.txt
          with open(
              os.path.join(data_path, _KERAS_SAVE_FORMAT_PATH), "w"
          ) as f:
              f.write(save_format)
  
          # save keras model
          # To maintain prior behavior, when the format is HDF5, we save
          # with the h5 file extension. Otherwise, model_path is a directory
          # where the saved_model.pb will be stored (for SavedModel format)
          # For tensorflow 2.16.0 (including dev version),
          # it only supports saving model in .h5 or .keras format
          if save_format == "h5":
              file_extension = ".h5"
          elif Version(tf.__version__).release >= (2, 16):
              file_extension = ".keras"
          else:
              file_extension = ""
          model_path = (
              os.path.join(path, model_subpath) + file_extension
          )
          if path.startswith("/dbfs/"):
              # The Databricks Filesystem uses a FUSE implementation that does not support
              # random writes. It causes an error.
              with tempfile.NamedTemporaryFile(suffix=".h5") as f:
                  model.save(f.name, **keras_model_kwargs)
                  f.flush()  # force flush the data
                  shutil.copy2(src=f.name, dst=model_path)
          else:
              model.save(model_path, **keras_model_kwargs)
  
          pyfunc_options = {
              "data": data_subpath,
          }
  
          flavor_options = {
              **pyfunc_options,
              "model_type": _MODEL_TYPE_KERAS,
              "keras_version": tf.__version__,
              "save_format": save_format,
          }
      elif isinstance(model, tf.Module):
          saved_model_kwargs = saved_model_kwargs or {}
          model_dir_subpath = "tf2model"
          model_path = os.path.join(path, model_dir_subpath)
          tf.saved_model.save(model, model_path, **saved_model_kwargs)
          pyfunc_options = {}
          flavor_options = {
              "saved_model_dir": model_dir_subpath,
              "model_type": _MODEL_TYPE_TF2_MODULE,
          }
      else:
          raise MlflowException(f"Unknown model type: {type(model)}")
  
      # update flavor info to mlflow_model
      mlflow_model.add_flavor(
          FLAVOR_NAME, code=code_dir_subpath, **flavor_options
      )
  
      # append loader_module, data and env data to mlflow_model
      pyfunc.add_to_model(
          mlflow_model,
          loader_module="mlflow.tensorflow",
          conda_env=_CONDA_ENV_FILE_NAME,
          python_env=_PYTHON_ENV_FILE_NAME,
          code=code_dir_subpath,
          **pyfunc_options,
      )
  
      # add model file size to mlflow_model
      if size := get_total_file_size(path):
          mlflow_model.model_size_bytes = size
  
      # save mlflow_model to path/MLmodel
      mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
  
      include_cloudpickle = (
          custom_objects is not None
          or get_global_custom_objects() is not None
      )
      if conda_env is None:
          if pip_requirements is None:
              default_reqs = get_default_pip_requirements(
                  include_cloudpickle
              )
              # To ensure `_load_pyfunc` can successfully load the model during the dependency
              # inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
              inferred_reqs = mlflow.models.infer_pip_requirements(
                  path, FLAVOR_NAME, fallback=default_reqs
              )
              default_reqs = sorted(
                  set(inferred_reqs).union(default_reqs)
              )
          else:
              default_reqs = None
          (
              conda_env,
              pip_requirements,
              pip_constraints,
          ) = _process_pip_requirements(
              default_reqs,
              pip_requirements,
              extra_pip_requirements,
          )
      else:
          (
              conda_env,
              pip_requirements,
              pip_constraints,
          ) = _process_conda_env(conda_env)
  
      with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
          yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
  
      # Save `constraints.txt` if necessary
      if pip_constraints:
          write_to(
              os.path.join(path, _CONSTRAINTS_FILE_NAME),
              "\n".join(pip_constraints),
          )
  
      # Save `requirements.txt`
      write_to(
          os.path.join(path, _REQUIREMENTS_FILE_NAME),
          "\n".join(pip_requirements),
      )
  
      _PythonEnv.current().to_yaml(
          os.path.join(path, _PYTHON_ENV_FILE_NAME)
      )
  
imports: import atexit
  import importlib
  import logging
  import os
  import shutil
  import tempfile
  import warnings
  from typing import Any, Dict, NamedTuple, Optional
  import numpy
  import pandas
  import yaml
  from packaging.version import Version
  import mlflow
  from mlflow import pyfunc
  from mlflow.data.code_dataset_source import CodeDatasetSource
  from mlflow.data.numpy_dataset import from_numpy
  from mlflow.data.tensorflow_dataset import from_tensorflow
  from mlflow.exceptions import INVALID_PARAMETER_VALUE, MlflowException
  from mlflow.models import Model, ModelInputExample, ModelSignature, infer_signature
  from mlflow.models.model import MLMODEL_FILE_NAME
  from mlflow.models.signature import _infer_signature_from_input_example
  from mlflow.models.utils import _save_example
  from mlflow.tensorflow.callback import MLflowCallback
  from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
  from mlflow.tracking.artifact_utils import _download_artifact_from_uri
  from mlflow.tracking.context import registry
  from mlflow.types.schema import TensorSpec
  from mlflow.utils import is_iterator
  from mlflow.utils.autologging_utils import PatchFunction, autologging_integration, batch_metrics_logger, get_autologging_config, log_fn_args_as_params, picklable_exception_safe_function, resolve_input_example_and_signature, safe_patch
  from mlflow.utils.autologging_utils.metrics_queue import add_to_metrics_queue, flush_metrics_queue
  from mlflow.utils.docstring_utils import LOG_MODEL_PARAM_DOCS, format_docstring
  from mlflow.utils.environment import _CONDA_ENV_FILE_NAME, _CONSTRAINTS_FILE_NAME, _PYTHON_ENV_FILE_NAME, _REQUIREMENTS_FILE_NAME, _mlflow_conda_env, _process_conda_env, _process_pip_requirements, _PythonEnv, _validate_env_arguments
  from mlflow.utils.file_utils import get_total_file_size, write_to
  from mlflow.utils.model_utils import _add_code_from_conf_to_system_path, _get_flavor_configuration, _validate_and_copy_code_paths, _validate_and_prepare_target_save_path
  from mlflow.utils.requirements_utils import _get_pinned_requirement
  from mlflow.utils.time import get_current_time_millis
label_desc: Use image node_id 6 for calling a global function with example usage: save_model(model, path, conda_env, code_paths, mlflow_model, custom_objects, signature, input_example, pip_requirements, extra_pip_requirements, saved_model_kwargs, keras_model_kwargs, metadata) without return types