def evaluate_quadratic(J, g, s, diag=None):
    """Compute values of a quadratic function arising in least squares.
    The function is 0.5 * s.T * (J.T * J + diag) * s + g.T * s.
    """
    if s.dim() == 1:
        Js = J.mv(s)
        q = Js.dot(Js)
        if diag is not None:
            q += s.dot(s * diag)
    else:
        Js = J.matmul(s.T)
        q = Js.square().sum(0)
        if diag is not None:
            q += (diag * s.square()).sum(1)

    l = s.matmul(g)

    return 0.5 * q + l