code: def backward(heads, head_grads=None, retain_graph=False, train_mode=True): #pylint: disable=redefined-outer-name
      
      head_handles, hgrad_handles = _parse_head(heads, head_grads)
  
      check_call(_LIB.MXAutogradBackwardEx(
          len(head_handles),
          head_handles,
          hgrad_handles,
          0,
          ctypes.c_void_p(0),
          ctypes.c_int(retain_graph),
          ctypes.c_int(0),
          ctypes.c_int(train_mode),
          ctypes.c_void_p(0),
          ctypes.c_void_p(0)))
docs: Compute the gradients of heads w.r.t previously marked variables.
  
      Parameters
      ----------
      heads: NDArray or list of NDArray
          Output NDArray(s)
      head_grads: NDArray or list of NDArray or None
          Gradients with respect to heads.
      train_mode: bool, optional
          Whether to do backward for training or predicting.