code: def create_default_update_fn(optimizer, model_loss):
      """
      This function calls the update function, to implement the backpropogation
      """
  
      @jax.jit
      def update(
          params, opt_state, batch, target, weights, rng
      ) -> Tuple[hk.Params, optax.OptState, jnp.ndarray]:
          batch_loss, grads = jax.value_and_grad(model_loss)(
              params, batch, target, weights, rng
          )
          updates, opt_state = optimizer.update(grads, opt_state)
          new_params = optax.apply_updates(params, updates)
          return new_params, opt_state, batch_loss
  
      return update
  
imports: import numpy
  import time
  import logging
  from collections.abc import Sequence
  from deepchem.data import Dataset, NumpyDataset
  from deepchem.metrics import Metric
  from deepchem.models.models import Model
  from deepchem.models.losses import Loss
  from deepchem.models.optimizers import Optimizer, Adam
  from deepchem.utils.evaluate import GeneratorEvaluator
  from deepchem.trans.transformers import Transformer, undo_transforms
  from typing import Any, Callable, Iterable, List, Optional, Tuple, Union, Sequence
  from deepchem.utils.typing import LossFn, OneOrMany
  import jax.numpy
  import jax
  import haiku
  import optax
  import warnings
label_desc: Use image node_id 2 for calling a global function with example usage: create_default_update_fn(optimizer, model_loss) and returns: update, new_params,  opt_state,  batch_loss