Compute the gradients of heads w.r.t variables. Gradients will be
    returned as new NDArrays instead of stored into `variable.grad`.
    Supports recording gradient graph for computing higher order gradients.

    .. note::

      Currently only a very limited set of operators support higher order \
      gradients.

    Parameters
    ----------
    heads: NDArray or list of NDArray
        Output NDArray(s)
    variables: NDArray or list of NDArray
        Input variables to compute gradients for.
    head_grads: NDArray or list of NDArray or None
        Gradients with respect to heads.
    retain_graph: bool
        Whether to keep computation graph to differentiate again, instead
        of clearing history and release memory. Defaults to the same value
        as create_graph.
    create_graph: bool
        Whether to record gradient graph for computing higher order
    train_mode: bool, optional
        Whether to do backward for training or prediction.

    Returns
    -------
    NDArray or list of NDArray:
        Gradients with respect to variables.

    Examples
    --------
    >>> x = mx.nd.ones((1,))
    >>> x.attach_grad()
    >>> with mx.autograd.record():
    ...     z = mx.nd.elemwise_add(mx.nd.exp(x), x)
    >>> dx = mx.autograd.grad(z, [x], create_graph=True)
    >>> print(dx)
    [
    [ 3.71828175]
    <NDArray 1 @cpu(0)>]

def grad(heads, variables, head_grads=None, retain_graph=None, create_graph=False,
         train_mode=True):  #pylint: disable=redefined-outer-name
    """Compute the gradients of heads w.r.t variables. Gradients will be
    returned as new NDArrays instead of stored into `variable.grad`.
    Supports recording gradient graph for computing higher order gradients.

    .. note::

      Currently only a very limited set of operators support higher order \
      gradients.

    Parameters
    ----------
    heads: NDArray or list of NDArray
        Output NDArray(s)
    variables: NDArray or list of NDArray
        Input variables to compute gradients for.
    head_grads: NDArray or list of NDArray or None
        Gradients with respect to heads.
    retain_graph: bool
        Whether to keep computation graph to differentiate again, instead
        of clearing history and release memory. Defaults to the same value
        as create_graph.
    create_graph: bool
        Whether to record gradient graph for computing higher order
    train_mode: bool, optional
        Whether to do backward for training or prediction.

    Returns
    -------
    NDArray or list of NDArray:
        Gradients with respect to variables.

    Examples
    --------
    >>> x = mx.nd.ones((1,))
    >>> x.attach_grad()
    >>> with mx.autograd.record():
    ...     z = mx.nd.elemwise_add(mx.nd.exp(x), x)
    >>> dx = mx.autograd.grad(z, [x], create_graph=True)
    >>> print(dx)
    [
    [ 3.71828175]
    <NDArray 1 @cpu(0)>]
    """
    head_handles, hgrad_handles = _parse_head(heads, head_grads)

    if isinstance(variables, NDArray):
        variables = [variables]
    else:
        assert len(variables), "variables cannot be an empty list."
    var_handles = c_handle_array(variables)

    retain_graph = retain_graph if retain_graph is not None else create_graph
    grad_vars = ctypes.POINTER(NDArrayHandle)()
    grad_stypes = ctypes.POINTER(ctypes.c_int)()

    check_call(_LIB.MXAutogradBackwardEx(
        len(head_handles),
        head_handles,
        hgrad_handles,
        len(var_handles),
        var_handles,
        ctypes.c_int(retain_graph),
        ctypes.c_int(create_graph),
        ctypes.c_int(train_mode),
        ctypes.byref(grad_vars),
        ctypes.byref(grad_stypes)))

    ret = [_ndarray_cls(ctypes.cast(grad_vars[i], NDArrayHandle),
                        stype=grad_stypes[i])
           for i in range(len(var_handles))]
    if isinstance(variables, NDArray):
        return ret[0]
    return ret