code: def update_model(state, grads):
      return state.apply_gradients(grads=grads)
  
file_path: flax/examples/mnist/train.py
imports: from absl import logging
  from flax import linen
  from flax.metrics import tensorboard
  from flax.training import train_state
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: update_model(state, grads) and returns: state