code: def autolog(
      every_n_iter=1,
      log_models=True,
      log_datasets=True,
      disable=False,
      exclusive=False,
      disable_for_unsupported_versions=False,
      silent=False,
      registered_model_name=None,
      log_input_examples=False,
      log_model_signatures=True,
      saved_model_kwargs=None,
      keras_model_kwargs=None,
      extra_tags=None,
  ):  # pylint: disable=unused-argument
      # pylint: disable=no-name-in-module
      """
      Enables autologging for ``tf.keras``.
      Note that only ``tensorflow>=2.3`` are supported.
      As an example, try running the
      `Keras/TensorFlow example <https://github.com/mlflow/mlflow/blob/master/examples/keras/train.py>`_.
  
      For each TensorFlow module, autologging captures the following information:
  
      **tf.keras**
       - **Metrics** and **Parameters**
  
        - Training and validation loss.
        - User-specified metrics.
        - Optimizer config, e.g., learning_rate, momentum, etc.
        - Training configs, e.g., epochs, batch_size, etc.
  
       - **Artifacts**
  
        - Model summary on training start.
        - Saved Keras model in `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ format.
        - TensorBoard logs on training end.
  
      **tf.keras.callbacks.EarlyStopping**
       - **Metrics** and **Parameters**
  
        - Metrics from the ``EarlyStopping`` callbacks: ``stopped_epoch``, ``restored_epoch``,
          ``restore_best_weight``, etc
        - ``fit()`` or ``fit_generator()`` parameters associated with ``EarlyStopping``:
          ``min_delta``, ``patience``, ``baseline``, ``restore_best_weights``, etc
  
      Refer to the autologging tracking documentation for more
      information on `TensorFlow workflows
      <https://www.mlflow.org/docs/latest/tracking.html#tensorflow-and-keras-experimental>`_.
  
      Note that autologging cannot be used together with explicit MLflow callback, i.e.,
      `mlflow.tensorflow.MLflowCallback`, because it will cause the same metrics to be logged twice.
      If you want to include `mlflow.tensorflow.MLflowCallback` in the callback list, please turn off
      autologging by calling `mlflow.tensorflow.autolog(disable=True)`.
  
      Args:
          every_n_iter: The frequency with which metrics should be logged. For example, a value of
              100 will log metrics at step 0, 100, 200, etc.
          log_models: If ``True``, trained models are logged as MLflow model artifacts.
              If ``False``, trained models are not logged.
          log_datasets: If ``True``, dataset information is logged to MLflow Tracking.
              If ``False``, dataset information is not logged.
          disable: If ``True``, disables the TensorFlow autologging integration. If ``False``,
              enables the TensorFlow integration autologging integration.
          exclusive: If ``True``, autologged content is not logged to user-created fluent runs.
              If ``False``, autologged content is logged to the active fluent run,
              which may be user-created.
          disable_for_unsupported_versions: If ``True``, disable autologging for versions of
              tensorflow that have not been tested against this version of the MLflow
              client or are incompatible.
          silent: If ``True``, suppress all event logs and warnings from MLflow during TensorFlow
              autologging. If ``False``, show all events and warnings during TensorFlow
              autologging.
          registered_model_name: If given, each time a model is trained, it is registered as a
              new model version of the registered model with this name.
              The registered model is created if it does not already exist.
          log_input_examples: If ``True``, input examples from training datasets are collected and
              logged along with tf/keras model artifacts during training. If
              ``False``, input examples are not logged.
          log_model_signatures: If ``True``,
              :py:class:`ModelSignatures <mlflow.models.ModelSignature>`
              describing model inputs and outputs are collected and logged along
              with tf/keras model artifacts during training. If ``False``,
              signatures are not logged. Note that logging TensorFlow models
              with signatures changes their pyfunc inference behavior when
              Pandas DataFrames are passed to ``predict()``.
              When a signature is present, an ``np.ndarray``
              (for single-output models) or a mapping from
              ``str`` -> ``np.ndarray`` (for multi-output models) is returned;
              when a signature is not present, a Pandas DataFrame is returned.
          saved_model_kwargs: a dict of kwargs to pass to ``tensorflow.saved_model.save`` method.
          keras_model_kwargs: a dict of kwargs to pass to ``keras_model.save`` method.
          extra_tags: A dictionary of extra tags to set on each managed run created by autologging.
      """
      import tensorflow as tf
  
      global _LOG_EVERY_N_STEPS
      _LOG_EVERY_N_STEPS = every_n_iter
  
      atexit.register(flush_metrics_queue)
  
      if Version(tf.__version__) < Version("2.3"):
          warnings.warn(
              "Could not log to MLflow. TensorFlow versions below 2.3 are not supported."
          )
          return
  
      @picklable_exception_safe_function
      def _get_early_stop_callback(callbacks):
          for callback in callbacks:
              if isinstance(callback, tf.keras.callbacks.EarlyStopping):
                  return callback
          return None
  
      def _log_early_stop_callback_params(callback):
          if callback:
              try:
                  earlystopping_params = {
                      "monitor": callback.monitor,
                      "min_delta": callback.min_delta,
                      "patience": callback.patience,
                      "baseline": callback.baseline,
                      "restore_best_weights": callback.restore_best_weights,
                  }
                  mlflow.log_params(earlystopping_params)
              except Exception:
                  return
  
      def _get_early_stop_callback_attrs(callback):
          try:
              return (
                  callback.stopped_epoch,
                  callback.restore_best_weights,
                  callback.patience,
              )
          except Exception:
              return None
  
      def _log_early_stop_callback_metrics(
          callback, history, metrics_logger
      ):
          if callback is None or not callback.model.stop_training:
              return
  
          callback_attrs = _get_early_stop_callback_attrs(callback)
          if callback_attrs is None:
              return
  
          stopped_epoch, restore_best_weights, _ = callback_attrs
          metrics_logger.record_metrics(
              {"stopped_epoch": stopped_epoch}
          )
  
          if not restore_best_weights or callback.best_weights is None:
              return
  
          monitored_metric = history.history.get(callback.monitor)
          if not monitored_metric:
              return
  
          initial_epoch = history.epoch[0]
          # If `monitored_metric` contains multiple best values (e.g. [0.1, 0.1, 0.2] where 0.1 is
          # the minimum loss), the epoch corresponding to the first occurrence of the best value is
          # the best epoch. In keras > 2.6.0, the best epoch can be obtained via the `best_epoch`
          # attribute of an `EarlyStopping` instance: https://github.com/keras-team/keras/pull/15197
          restored_epoch = initial_epoch + monitored_metric.index(
              callback.best
          )
          metrics_logger.record_metrics(
              {"restored_epoch": restored_epoch}
          )
          restored_index = history.epoch.index(restored_epoch)
          restored_metrics = {
              key: metrics[restored_index]
              for key, metrics in history.history.items()
          }
          # Checking that a metric history exists
          metric_key = next(iter(history.history), None)
          if metric_key is not None:
              metrics_logger.record_metrics(
                  restored_metrics, stopped_epoch + 1
              )
  
      def _log_keras_model(history, args):
          def _infer_model_signature(input_data_slice):
              # In certain TensorFlow versions, calling `predict()` on model may modify
              # the `stop_training` attribute, so we save and restore it accordingly
              original_stop_training = history.model.stop_training
              model_output = history.model.predict(input_data_slice)
              history.model.stop_training = original_stop_training
              return infer_signature(input_data_slice, model_output)
  
          from mlflow.tensorflow._autolog import (
              extract_tf_keras_input_example,
          )
  
          def _get_tf_keras_input_example_slice():
              input_training_data = args[0]
              keras_input_example_slice = (
                  extract_tf_keras_input_example(input_training_data)
              )
              if keras_input_example_slice is None:
                  raise MlflowException(
                      "Cannot log input example or model signature for input with type"
                      f" {type(input_training_data)}. TensorFlow Keras autologging can"
                      " only log input examples and model signatures for the following"
                      " input types: numpy.ndarray, dict[string -> numpy.ndarray],"
                      " tensorflow.keras.utils.Sequence, and"
                      " tensorflow.data.Dataset (TensorFlow >= 2.1.0 required)",
                      INVALID_PARAMETER_VALUE,
                  )
              return keras_input_example_slice
  
          (
              input_example,
              signature,
          ) = resolve_input_example_and_signature(
              _get_tf_keras_input_example_slice,
              _infer_model_signature,
              log_input_examples,
              log_model_signatures,
              _logger,
          )
  
          log_model(
              model=history.model,
              artifact_path="model",
              input_example=input_example,
              signature=signature,
              registered_model_name=get_autologging_config(
                  FLAVOR_NAME, "registered_model_name", None
              ),
              saved_model_kwargs=saved_model_kwargs,
              keras_model_kwargs=keras_model_kwargs,
          )
  
      class FitPatch(PatchFunction):
          def __init__(self):
              self.log_dir = None
  
          def _patch_implementation(
              self, original, inst, *args, **kwargs
          ):  # pylint: disable=arguments-differ
              unlogged_params = [
                  "self",
                  "x",
                  "y",
                  "callbacks",
                  "validation_data",
                  "verbose",
              ]
  
              batch_size = None
              try:
                  is_single_input_model = isinstance(
                      inst.input_shape, tuple
                  )
                  training_data = (
                      kwargs["x"] if "x" in kwargs else args[0]
                  )
                  if isinstance(
                      training_data, tf.data.Dataset
                  ) and hasattr(training_data, "_batch_size"):
                      batch_size = training_data._batch_size.numpy()
                  elif isinstance(
                      training_data, tf.keras.utils.Sequence
                  ):
                      first_batch_inputs, *_ = training_data[0]
                      if is_single_input_model:
                          batch_size = len(first_batch_inputs)
                      else:
                          batch_size = len(first_batch_inputs[0])
                  elif is_iterator(training_data):
                      peek = next(training_data)
                      batch_size = (
                          len(peek[0])
                          if is_single_input_model
                          else len(peek[0][0])
                      )
  
                      def __restore_generator(prev_generator):
                          yield peek
                          yield from prev_generator
  
                      restored_generator = __restore_generator(
                          training_data
                      )
                      if "x" in kwargs:
                          kwargs["x"] = restored_generator
                      else:
                          args = (restored_generator,) + args[1:]
              except Exception as e:
                  _logger.warning(
                      "Encountered unexpected error while inferring batch size from training"
                      " dataset: %s",
                      e,
                  )
  
              if batch_size is not None:
                  mlflow.log_param("batch_size", batch_size)
                  unlogged_params.append("batch_size")
  
              log_fn_args_as_params(
                  original, args, kwargs, unlogged_params
              )
  
              run_id = mlflow.active_run().info.run_id
              with batch_metrics_logger(run_id) as metrics_logger:
                  # Check if the 'callback' argument of fit() is set positionally
                  if len(args) >= 6:
                      # Convert the positional training function arguments to a list in order to
                      # mutate the contents
                      args = list(args)
                      # Make a shallow copy of the preexisting callbacks to avoid permanently
                      # modifying their contents for future training invocations. Introduce
                      # TensorBoard & tf.keras callbacks if necessary
                      callbacks = list(args[5])
                      callbacks, self.log_dir = _setup_callbacks(
                          callbacks, metrics_logger
                      )
                      # Replace the callbacks positional entry in the copied arguments and convert
                      # the arguments back to tuple form for usage in the training function
                      args[5] = callbacks
                      args = tuple(args)
                  else:
                      # Make a shallow copy of the preexisting callbacks and introduce TensorBoard
                      # & tf.keras callbacks if necessary
                      callbacks = list(kwargs.get("callbacks") or [])
                      (
                          kwargs["callbacks"],
                          self.log_dir,
                      ) = _setup_callbacks(callbacks, metrics_logger)
  
                  early_stop_callback = _get_early_stop_callback(
                      callbacks
                  )
                  _log_early_stop_callback_params(early_stop_callback)
  
                  if log_datasets:
                      try:
                          context_tags = context_registry.resolve_tags()
                          source = CodeDatasetSource(tags=context_tags)
  
                          x = kwargs["x"] if "x" in kwargs else args[0]
                          if "y" in kwargs:
                              y = kwargs["y"]
                          elif len(args) >= 2:
                              y = args[1]
                          else:
                              y = None
  
                          if "validation_data" in kwargs:
                              validation_data = kwargs[
                                  "validation_data"
                              ]
                          elif len(args) >= 8:
                              validation_data = args[7]
                          else:
                              validation_data = None
                          _log_tensorflow_dataset(
                              x, source, "train", targets=y
                          )
                          if validation_data is not None:
                              _log_tensorflow_dataset(
                                  validation_data, source, "eval"
                              )
  
                      except Exception as e:
                          _logger.warning(
                              "Failed to log training dataset information to "
                              "MLflow Tracking. Reason: %s",
                              e,
                          )
  
                  history = original(inst, *args, **kwargs)
  
                  if log_models:
                      _log_keras_model(history, args)
  
                  _log_early_stop_callback_metrics(
                      callback=early_stop_callback,
                      history=history,
                      metrics_logger=metrics_logger,
                  )
  
                  flush_metrics_queue()
                  mlflow.log_artifacts(
                      local_dir=self.log_dir.location,
                      artifact_path="tensorboard_logs",
                  )
              if self.log_dir.is_temp:
                  shutil.rmtree(self.log_dir.location)
              return history
  
          def _on_exception(self, exception):
              if (
                  self.log_dir is not None
                  and self.log_dir.is_temp
                  and os.path.exists(self.log_dir.location)
              ):
                  shutil.rmtree(self.log_dir.location)
  
      managed = [
          (tf.keras.Model, "fit", FitPatch),
      ]
  
      for p in managed:
          safe_patch(
              FLAVOR_NAME, *p, manage_run=True, extra_tags=extra_tags
          )
  
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 18 for calling a global function with example usage: autolog(every_n_iter, log_models, log_datasets, disable, exclusive, disable_for_unsupported_versions, silent, registered_model_name, log_input_examples, log_model_signatures, saved_model_kwargs, keras_model_kwargs, extra_tags) and returns: None, callback,  callback,  callback, infer_signature, keras_input_example_slice, history, callback, None