code: def run_fn(fn_args: tfx.components.FnArgs):
      """Train the model based on given args.
  
      Args:
        fn_args: Holds args used to train the model as name/value pairs.
      """
      tf_transform_output = tft.TFTransformOutput(
          fn_args.transform_output
      )
  
      train_dataset = base.input_fn(
          fn_args.train_files,
          fn_args.data_accessor,
          tf_transform_output,
          40,
      )
      eval_dataset = base.input_fn(
          fn_args.eval_files,
          fn_args.data_accessor,
          tf_transform_output,
          40,
      )
  
      mirrored_strategy = tf.distribute.MirroredStrategy()
      with mirrored_strategy.scope():
          model = base.build_keras_model()
  
      # Write logs to path
      tensorboard_callback = tf.keras.callbacks.TensorBoard(
          log_dir=fn_args.model_run_dir, update_freq="epoch"
      )
  
      model.fit(
          train_dataset,
          steps_per_epoch=fn_args.train_steps,
          validation_data=eval_dataset,
          validation_steps=fn_args.eval_steps,
          callbacks=[tensorboard_callback],
      )
  
      signatures = {
          "serving_default": _get_serve_tf_examples_fn(
              model, tf_transform_output
          ).get_concrete_function(
              tf.TensorSpec(
                  shape=[None, 784],
                  dtype=tf.float32,
                  name="image_floats",
              )
          )
      }
      temp_saving_model_dir = os.path.join(
          fn_args.serving_model_dir, "temp"
      )
      model.save(
          temp_saving_model_dir, save_format="tf", signatures=signatures
      )
  
      tfrw = rewriter_factory.create_rewriter(
          rewriter_factory.TFLITE_REWRITER, name="tflite_rewriter"
      )
      converters.rewrite_saved_model(
          temp_saving_model_dir,
          fn_args.serving_model_dir,
          tfrw,
          rewriter.ModelType.TFLITE_MODEL,
      )
  
      tfx.dsl.io.fileio.rmtree(temp_saving_model_dir)
  
file_path: tfx/tfx/examples/mnist/mnist_utils_native_keras_lite.py
imports: import os
  import tensorflow
  import tensorflow_transform
  from tfx import v1
  from tfx.components.trainer.rewriting import converters
  from tfx.components.trainer.rewriting import rewriter
  from tfx.components.trainer.rewriting import rewriter_factory
  from tfx.examples.mnist import mnist_utils_native_keras_base
label_desc: Use image node_id 3 for calling a global function with example usage: run_fn(fn_args) without return types