code: def run_epoch(state, dataset, train=True):
      epoch_loss = []
      epoch_accuracy = []
  
      for inputs, labels in dataset:
          grads, loss, accuracy = apply_model(state, inputs, labels)
          if train:
              state = update_model(state, grads)
          epoch_loss.append(loss)
          epoch_accuracy.append(accuracy)
      loss = np.mean(epoch_loss)
      accuracy = np.mean(epoch_accuracy)
      return state, loss, accuracy
  
imports: from absl import app
  from absl import flags
  import os
  import time
  from flax import linen
  from flax.training import train_state
  import jax
  from jax import lax
  from jax import numpy
  import optax
  import numpy
  import tensorflow
  import tensorflowjs
  import input_pipeline
label_desc: Use image node_id 3 for calling a global function with example usage: run_epoch(state, dataset, train) and returns: state,  loss,  accuracy