code: def create_default_update_fn(
      optimizer: optax.GradientTransformation, model_loss: Callable
  ):
      """
      This function calls the update function, to implement the backpropagation
      """
  
      @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, target, weights, rng, *batch
          )
          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
  from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
  from collections.abc import Sequence
  from deepchem.data import Dataset
  from deepchem.models.losses import Loss
  from deepchem.models.optimizers import Optimizer
  from deepchem.utils.typing import LossFn, OneOrMany
  from deepchem.trans.transformers import Transformer, undo_transforms
  import jax
  import jax.numpy
  import haiku
  import optax
  from deepchem.models.jax_models.jax_model import JaxModel
  from deepchem.models.jax_models.jax_model import create_default_gradient_fn, create_default_eval_fn
  import logging
  import warnings
label_desc: Use image node_id 1 for calling a global function with example usage: create_default_update_fn(optimizer, model_loss) and returns: update, new_params,  opt_state,  batch_loss