import torch


def compute_damped_gn_update(jacobian, output_error, damping):
    """
    Compute the damped Gauss-Newton update, based on the given jacobian and
    output error.
    Args:
        jacobian (torch.Tensor): 2D tensor containing the Jacobian of the
            flattened output with respect to the flattened parameters for which
            the GN update is computed.
        output_error (torch.Tensor): tensor containing the gradient of the loss
            with respect to the output layer of the network.
        damping (float): positive damping hyperparameter

    Returns: the damped Gauss-Newton update for the parameters for which the
        jacobian was computed.

    """
    if damping < 0:
        raise ValueError('Positive value for damping expected, got '
                         '{}'.format(damping))
    # The jacobian also flattens the  output dimension, so we need to do
    # the same.
    output_error = output_error.view(-1, 1).detach()

    if damping == 0:
        # if the damping is 0, the curvature matrix C=J^TJ can be
        # rank deficit. Therefore, it is numerically best to compute the
        # pseudo inverse explicitly and after that multiply with it.
        jacobian_pinv = torch.pinverse(jacobian)
        gn_updates = jacobian_pinv.mm(output_error)
    else:
        # If damping is greater than 0, the curvature matrix C will be
        # positive definite and symmetric. Numerically, it is the most
        # efficient to use the cholesky decomposition to compute the
        # resulting Gauss-newton updates

        # As (J^T*J + l*I)^{-1}*J^T = J^T*(JJ^T + l*I)^{-1}, we select
        # the one which is most computationally efficient, depending on
        # the number of rows and columns of J (we want to take the inverse
        # of the smallest possible matrix, as this is the most expensive
        # operation. Note that we solve a linear system with cholesky
        # instead of explicitly computing the inverse, as this is more
        # efficient.
        if jacobian.shape[0] >= jacobian.shape[1]:
            G = jacobian.t().mm(jacobian)
            C = G + damping * torch.eye(G.shape[0])
            C_cholesky = torch.cholesky(C)
            jacobian_error = jacobian.t().matmul(output_error)
            gn_updates = torch.cholesky_solve(jacobian_error, C_cholesky)
        else:
            G = jacobian.mm(jacobian.t())
            C = G + damping * torch.eye(G.shape[0])
            C_cholesky = torch.cholesky(C)
            inverse_error = torch.cholesky_solve(output_error, C_cholesky)
            gn_updates = jacobian.t().matmul(inverse_error)

    return gn_updates