code: def model_main(learning_rate=0.01, epochs=3, **kwargs):
      """Train a model against the MNIST dataset and print performance metrics."""
      mnist = tf.keras.datasets.mnist
  
      (x_train, y_train), (x_test, y_test) = mnist.load_data()
      x_train, x_test = x_train / 255.0, x_test / 255.0
  
      model = get_keras_model()
  
      loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
          from_logits=True
      )
      optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
  
      model.compile(
          optimizer=optimizer, loss=loss_fn, metrics=["accuracy"]
      )
  
      # this is the mutable map where the MemoryReporter will keep its data.
      metrics = {}
      base = build_reporters(metrics)
  
      with mlflow.start_run():
          mlflow.log_params(
              {
                  **kwargs,
                  **{"learning_rate": learning_rate, "epochs": epochs},
              }
          )
          print(
              f"Training model with learning rate={learning_rate} for {epochs} epochs."
          )
          model.fit(x_train, y_train, epochs=epochs)
  
          print("Model performance: ")
          score, accuracy = model.evaluate(x_test, y_test)
          mlflow.log_params(
              {"final_score": score, "accuracy": accuracy}
          )
  
file_path: caliban/tutorials/uv-metrics/mnist.py
imports: import warnings
  from typing import Dict
  import mlflow
  import mlflow.keras
  import tensorflow
  import uv.reporter
  import uv.util
  from absl import app
  from uv.mlflow.reporter import MLFlowReporter
  import cli
label_desc: Use image node_id 3 for calling a global function with example usage: model_main(learning_rate, epochs) without return types