def gradient_descent_with_momentum(loss_grad, theta, lr=0.001, momentum=0.9,
    maxsteps=100):
    """Gradient descent with momentum.

    Find an approximation of the vector theta that minimizes an
    objective function.

    Parameters
    ----------
    loss_grad : callable
        function computing the gradient of the objective.
    theta : ndarray, shape (n,)
        initial value.
    lr : float, optional
        learning rate.
    momentum : float, optional
        momentum term.
    maxsteps : int, optional
        maximum number of iterations.

    Returns
    -------
    ndarray, shape (n,)
        the final estimate found by the algorithm.

    """
    v = theta.copy()
    for _ in range(maxsteps):
        delta = momentum * v + loss_grad(theta)
        theta -= lr * delta
    return theta
