def lu_factor(A, rho, check_finite=True):
    r"""
    Compute LU factorisation of either :math:`A^T A + \rho I` or
    :math:`A A^T + \rho I`, depending on which matrix is smaller.

    Parameters
    ----------
    A : array_like
      Array :math:`A`
    rho : float
      Scalar :math:`\rho`
    check_finite : bool, optional (default False)
      Flag indicating whether the input array should be checked for Inf
      and NaN values

    Returns
    -------
    lu : ndarray
      Matrix containing U in its upper triangle, and L in its lower triangle,
      as returned by :func:`scipy.linalg.lu_factor`
    piv : ndarray
      Pivot indices representing the permutation matrix P, as returned by
      :func:`scipy.linalg.lu_factor`
    """

    N, M = A.shape
    # If N < M it is cheaper to factorise A*A^T + rho*I and then use the
    # matrix inversion lemma to compute the inverse of A^T*A + rho*I
    if N >= M:
        lu, piv = linalg.lu_factor(A.T.dot(A) +
                                   rho * np.identity(M, dtype=A.dtype),
                                   check_finite=check_finite)
    else:
        lu, piv = linalg.lu_factor(A.dot(A.T) +
                                   rho * np.identity(N, dtype=A.dtype),
                                   check_finite=check_finite)
    return lu, piv