code: def _make_trained_model(
      train_data: tf.data.Dataset,
      eval_data: tf.data.Dataset,
      num_epochs: int,
      steps_per_epoch: int,
      eval_steps_per_epoch: int,
      tensorboard_log_dir: str,
  ):
      """Execute model training and evaluation loop.
  
      Args:
        train_data: a dataset with training pairs (_InputBatch, _LabelBatch).
        eval_data: a dataset with evaluation pairs (_InputBatch, _LabelBatch).
        num_epochs: number of training epochs.
        steps_per_epoch: number of steps for a training epoch. Should be the number
           of samples in your train_data divided by the batch size.
        eval_steps_per_epoch: number of steps for evaluation at the end of each
           training epoch. Should be the number of samples in your eval_data
           divided by the batch size.
        tensorboard_log_dir: Directory where the tensorboard summaries are written.
  
      Returns:
        An instance of tf.Model.
      """
      learning_rate = 1e-2
  
      rng = jax.random.PRNGKey(0)
  
      summary_writer = tensorboard.SummaryWriter(tensorboard_log_dir)
      summary_writer.hparams(
          dict(
              learning_rate=learning_rate,
              num_epochs=num_epochs,
              steps_per_epoch=steps_per_epoch,
              eval_steps_per_epoch=eval_steps_per_epoch,
          )
      )
  
      rng, init_rng = jax.random.split(rng)
      # Initialize with some fake data of the proper shape.
      init_val = dict(
          (feature, jnp.array([[1.0]], dtype=jnp.float32))
          for feature in _FEATURE_KEYS_XF
      )
      model = _FlaxPenguinModel()
      params = model.init(init_rng, init_val)["params"]
  
      tx = optax.adam(learning_rate=learning_rate)
      opt_state = tx.init(params)
  
      for epoch in range(1, num_epochs + 1):
          params, opt_state, train_metrics = _train_epoch(
              model, tx, params, opt_state, train_data, steps_per_epoch
          )
          absl.logging.info(
              "Flax train epoch: %d, loss: %.4f, accuracy: %.2f",
              epoch,
              train_metrics["loss"],
              train_metrics["accuracy"] * 100,
          )
  
          eval_metrics = _eval_epoch(
              model, params, eval_data, eval_steps_per_epoch
          )
          absl.logging.info(
              "Flax eval epoch: %d, loss: %.4f, accuracy: %.2f",
              epoch,
              eval_metrics["loss"],
              eval_metrics["accuracy"] * 100,
          )
          summary_writer.scalar(
              "epoch_train_loss", train_metrics["loss"], epoch
          )
          summary_writer.scalar(
              "epoch_train_accuracy", train_metrics["accuracy"], epoch
          )
          summary_writer.scalar(
              "epoch_eval_loss", eval_metrics["loss"], epoch
          )
          summary_writer.scalar(
              "epoch_eval_accuracy", eval_metrics["accuracy"], epoch
          )
  
      summary_writer.flush()
  
      # The prediction function for the trained model
      def predict(params: _Params, inputs: _InputBatch):
          return model.apply({"params": params}, inputs)
  
      trained_params = params
  
      # Convert the prediction function to TF, with a variable batch dimension
      # for all inputs.
      tf_fn = jax2tf.convert(
          predict,
          with_gradient=False,
          enable_xla=True,
          polymorphic_shapes=(None, "(b, 1)"),
      )
  
      # Create tf.Variables for the parameters. If you want more useful variable
      # names, you can use `tree.map_structure_with_path` from the `dm-tree`
      # package.
      param_vars = tf.nest.map_structure(
          # Due to a bug in SavedModel it is not possible to use tf.GradientTape
          # on a function converted with jax2tf and loaded from SavedModel.
          # Thus, we mark the variables as non-trainable to ensure that users of
          # the SavedModel will not try to fine tune them.
          lambda param: tf.Variable(param, trainable=False),
          trained_params,
      )
      tf_graph = tf.function(
          lambda inputs: tf_fn(param_vars, inputs),
          autograph=False,
          experimental_compile=True,
      )
      return _SavedModelWrapper(tf_graph, param_vars)
  
imports: import functools
  from typing import Callable, Dict, List
  import absl
  from flax import linen
  from flax.metrics import tensorboard
  import jax
  from jax import numpy
  from jax.experimental import jax2tf
  import numpy
  import optax
  import tensorflow
  import tensorflow_transform
  from tfx import v1
  from tfx.examples.penguin import penguin_utils_base
label_desc: Use image node_id 1 for calling a global function with example usage: _make_trained_model(train_data, eval_data, num_epochs, steps_per_epoch, eval_steps_per_epoch, tensorboard_log_dir) and returns: _SavedModelWrapper, model